AskWikidata_Quickstart.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U23PvTMk2T2l"
      },
      "source": [
        "# AskWikidata\n",
        "\n",
        "\n",
        "A prototype for a Wikidata Question Answering System.\n",
        "https://github.com/rti/askwikidata\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t63hh1GTQSRL"
      },
      "source": [
        "## Quickstart Notebook"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/\n",
        "%rm -rf askwikidata"
      ],
      "metadata": {
        "id": "lXw03UtkFlFG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3TBYXNbo3s94",
        "scrolled": true,
        "collapsed": true
      },
      "outputs": [],
      "source": [
        "# Clone the askwikidata repo from Github.\n",
        "!git clone https://github.com/rti/askwikidata\n",
        "%cd /content/askwikidata"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%pip install \\\n",
        "  sentence-transformers==3.1.1 \\\n",
        "  langchain==0.3.0 \\\n",
        "  langchain-community==0.3.0 \\\n",
        "  annoy==1.17.3 \\\n",
        "  bitsandbytes==0.43.3"
      ],
      "metadata": {
        "collapsed": true,
        "id": "6f619pagRWLC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "iuGfAbH1f2a-"
      },
      "outputs": [],
      "source": [
        "# Unzip all cache files provided with the askwikidata repository.\n",
        "!cd /content/askwikidata && bunzip2 --force --keep *.bz2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "FXrXBcLU4Fz5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "bfba3c7d-3456-4482-f6a9-31e1a657a11e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "100% 13877/13877 [00:12<00:00, 1137.28it/s]\n"
          ]
        }
      ],
      "source": [
        "# Generate text representations of Wikidata items.\n",
        "!cd /content/askwikidata && python text_representation.py > text_representations.log"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x3UBuXzO8CAL",
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# Add askwikidata source to the python import paths.\n",
        "import sys\n",
        "sys.path.append('/content/askwikidata')\n",
        "\n",
        "# Change working to askwikidata\n",
        "import os\n",
        "os.chdir('/content/askwikidata')\n",
        "\n",
        "\n",
        "# Setup the actual AskWikidata RAG system.\n",
        "from askwikidata import AskWikidata\n",
        "\n",
        "config = {\n",
        "    \"chunk_size\": 1280,\n",
        "    \"chunk_overlap\": 0,\n",
        "    \"index_trees\": 1024,\n",
        "    \"retrieval_chunks\": 16,\n",
        "    \"context_chunks\": 5,\n",
        "    \"embedding_model_name\": \"BAAI/bge-small-en-v1.5\",\n",
        "    \"reranker_model_name\": \"BAAI/bge-reranker-base\",\n",
        "    \"qa_model_url\": \"Qwen/Qwen2.5-3B-Instruct\",\n",
        "}\n",
        "\n",
        "askwikidata = AskWikidata(**config)\n",
        "askwikidata.setup()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "SD5GEZjwFwR6",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e81e4e1c-3580-45d8-97d9-43f304adcbc4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Retrieving...\n",
            "Reranking...\n",
            "Generating...\n",
            "<|im_start|>assistant\n",
            "Kai Wegner since 2023-04-27 until today.<|im_end|>\n",
            "\n",
            "Sources:\n",
            "- https://www.wikidata.org/wiki/Q2079\n",
            "- https://www.wikidata.org/wiki/Q1757\n",
            "- https://www.wikidata.org/wiki/Q64\n",
            "- https://www.wikidata.org/wiki/Q1055\n"
          ]
        }
      ],
      "source": [
        "# Answer an example question.\n",
        "print(askwikidata.ask(\"Who is the current mayor of Berlin? Since when do they serve?\"))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}