rag_powered_text_to_sql_duckdb_phidata_llama_cpp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG-powered Text-To-SQL using PhiData, Ollama and Argilla"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this recipe, we will show how to create a powerful text-to-SQL engine that is powered by a RAG pipeline with [PhiData](https://docs.phidata.com/). We will use the [Ollama](https://github.com/ollama/ollama) to run the LLM and [DuckDB](https://duckdb.org/) powered to [run the SQL on top of datasets on the Hugging Face Hub](https://duckdb.org/2024/05/29/access-150k-plus-datasets-from-hugging-face-with-duckdb.html). Eventually, we will use Argilla to structure the data and review the textual prompts. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting started\n",
    "\n",
    "### Deploy the Argilla server¶\n",
    "\n",
    "If you already have deployed Argilla, you can skip this step. Otherwise, you can quickly deploy Argilla following [this guide](https://docs.argilla.io/latest/getting_started/quickstart/).\n",
    "\n",
    "### Install dependencies\n",
    "\n",
    "Since we will be using Ollama, we recommend installing ollama as [described on their GitHub repo](https://github.com/ollama/ollama). After that, we can deploy our ollama server and download and serve required model with.\n",
    "\n",
    "```bash\n",
    "ollama serve\n",
    "```\n",
    "\n",
    "```bash\n",
    "ollama run llama3.2\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Besides that, we will also be needing some Python dependencies. You can install them by running:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qqq bs4 phidata lancedb duckdb ollama datasets argilla "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create our RAG pipeline\n",
    "\n",
    "We want our model to have a better understanding of the DuckDB documentation and information. We will therefore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34mINFO    \u001b[0m Creating table: phi                                                    \n",
      "\u001b[34mINFO    \u001b[0m Creating table: phi                                                    \n",
      "\u001b[34mINFO    \u001b[0m Loading knowledge base                                                 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-10-08T14:07:21Z WARN  lance::dataset] No existing dataset at /tmp/lancedb/phi.lance, it will be created\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34mINFO    \u001b[0m Loaded \u001b[1;36m10\u001b[0m documents to knowledge base                                  \n",
      "\u001b[34mINFO    \u001b[0m Loaded \u001b[1;36m20\u001b[0m documents to knowledge base                                  \n",
      "\u001b[34mINFO    \u001b[0m Loaded \u001b[1;36m30\u001b[0m documents to knowledge base                                  \n",
      "\u001b[34mINFO    \u001b[0m Loaded \u001b[1;36m40\u001b[0m documents to knowledge base                                  \n",
      "\u001b[34mINFO    \u001b[0m Loaded \u001b[1;36m50\u001b[0m documents to knowledge base                                  \n"
     ]
    }
   ],
   "source": [
    "from phi.knowledge.website import WebsiteKnowledgeBase\n",
    "from phi.vectordb.lancedb import LanceDb\n",
    "from phi.embedder.ollama import OllamaEmbedder\n",
    "from phi.document.reader.website import WebsiteReader\n",
    "\n",
    "knowledge_base = WebsiteKnowledgeBase(\n",
    "    reader=WebsiteReader(chunk=False),\n",
    "    urls=[\n",
    "        \"https://duckdb.org/docs/sql/statements/overview\", \n",
    "        \"https://duckdb.org/docs/sql/data_types/overview\",\n",
    "        \"https://duckdb.org/docs/sql/expressions/overview\",\n",
    "        \"https://duckdb.org/docs/sql/functions/overview\",\n",
    "        \"https://duckdb.org/docs/sql/dialect/overview\"\n",
    "    ],\n",
    "    max_links=1000,\n",
    "    vector_db=LanceDb(\n",
    "        embedder=OllamaEmbedder(model=\"llama3.1\")\n",
    "    ),\n",
    "    max_depth=2\n",
    ")\n",
    "knowledge_base.load(recreate=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get dataset information\n",
    "\n",
    "Hugging Face datasets are hosted publicly but in order to generate synthetic text queries and their corresponding SQL, we will first need to obtain some informationf about this datasets. We will use this with the Hugging Face Hub API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List Hugging Face datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetInfo(id='google/frames-benchmark', author=None, sha=None, created_at=None, last_modified=None, private=None, gated=None, disabled=None, downloads=None, downloads_all_time=None, likes=None, paperswithcode_id=None, tags=None, trending_score=52, card_data=None, siblings=None)"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from huggingface_hub import HfApi\n",
    "\n",
    "datasets = list(HfApi().list_datasets(limit=100, gated=False, expand=[\"description\"]))\n",
    "datasets[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get dataset descriptions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Repo card metadata block was not found. Setting CardData to empty.\n",
      "Repo card metadata block was not found. Setting CardData to empty.\n",
      "Repo card metadata block was not found. Setting CardData to empty.\n",
      "Repo card metadata block was not found. Setting CardData to empty.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'# FRAMES: Factuality, Retrieval, And reasoning MEasurement Set\\n\\nFRAMES is a comprehensive evaluation dataset designed to test the capabilities of Retrieval-Augmented Generation (RAG) systems across factuality, retrieval accuracy, and reasoning.\\nOur paper with details and experiments is available on arXiv: [https://arxiv.org/abs/2409.12941](https://arxiv.org/abs/2409.12941).\\n\\n\\n## Dataset Overview\\n\\n- 824 challenging multi-hop questions requiring information from 2-15 Wikipedia articles\\n- Questions span diverse topics including history, sports, science, animals, health, etc.\\n- Each question is labeled with reasoning types: numerical, tabular, multiple constraints, temporal, and post-processing\\n- Gold answers and relevant Wikipedia articles provided for each question\\n\\n## Key Features\\n\\n- Tests end-to-end RAG capabilities in a unified framework\\n- Requires integration of information from multiple sources\\n- Incorporates complex reasoning and temporal disambiguation\\n- Designed to be challenging for state-of-the-art language models\\n\\n## Usage\\n\\nThis dataset can be used to:\\n- Evaluate RAG system performance \\n- Benchmark language model factuality and reasoning\\n- Develop and test multi-hop retrieval strategies\\n\\n## Baseline Results\\n\\nWe provide baseline results using state-of-the-art models like Gemini-Pro-1.5-0514:\\n\\n- Naive prompting: 40.8% accuracy\\n- BM25 retrieval (4 docs): 47.4% accuracy  \\n- Oracle retrieval: 72.9% accuracy\\n- Multi-step retrieval & reasoning: 66% accuracy\\n\\n## Citation\\n\\nIf you use this dataset in your research, please cite our paper:\\n\\n```\\n@misc{krishna2024factfetchreasonunified,\\n      title={Fact, Fetch, and Reason: A Unified Evaluation of Retrieval-Augmented Generation}, \\n      author={Satyapriya Krishna and Kalpesh Krishna and Anhad Mohananey and Steven Schwarcz and Adam Stambler and Shyam Upadhyay and Manaal Faruqui},\\n      year={2024},\\n      eprint={2409.12941},\\n      archivePrefix={arXiv},\\n      primaryClass={cs.CL},\\n      url={https://arxiv.org/abs/2409.12941}, \\n}\\n```\\n\\nWe hope FRAMES will be useful for advancing RAG systems and language model capabilities. For more details, please refer to our full paper.'"
      ]
     },
     "execution_count": 155,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from huggingface_hub import ModelCard\n",
    "\n",
    "card_content = []\n",
    "for dataset in datasets:\n",
    "    try:\n",
    "        card_content.append(\"---\".join(ModelCard.load(dataset.id, repo_type=\"dataset\").content.split(\"---\")[2:]).strip())\n",
    "    except Exception:\n",
    "        if hasattr(dataset, \"description\"):\n",
    "            card_content.append(dataset.description)\n",
    "        else:\n",
    "            card_content.append(None)\n",
    "card_content[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get dataset info from the datasets-server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('google/frames-benchmark',\n",
       " 'default',\n",
       " 'test',\n",
       " 824,\n",
       " {'Unnamed: 0': {'dtype': 'int64', '_type': 'Value'},\n",
       "  'Prompt': {'dtype': 'string', '_type': 'Value'},\n",
       "  'Answer': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_1': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_2': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_3': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_4': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_5': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_6': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_7': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_8': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_9': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_10': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wikipedia_link_11+': {'dtype': 'string', '_type': 'Value'},\n",
       "  'reasoning_types': {'dtype': 'string', '_type': 'Value'},\n",
       "  'wiki_links': {'dtype': 'string', '_type': 'Value'}})"
      ]
     },
     "execution_count": 173,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import requests\n",
    "\n",
    "def get_dataset_info(dataset_id):\n",
    "    url = f\"https://datasets-server.huggingface.co/info?dataset={dataset_id}\"\n",
    "    response = requests.get(url)\n",
    "    data = response.json()\n",
    "    if data.get(\"dataset_info\"):\n",
    "        dataset_info = data[\"dataset_info\"]\n",
    "        dataset_config = list(dataset_info.keys())[0]\n",
    "        dataset_split = list(dataset_info[dataset_config][\"splits\"].keys())[0]\n",
    "        dataset_size = dataset_info[dataset_config][\"splits\"][dataset_split][\"num_examples\"]\n",
    "        dataset_features = dataset_info[dataset_config][\"features\"]\n",
    "        return dataset_id, dataset_config, dataset_split, dataset_size, dataset_features\n",
    "    else:\n",
    "        return None, None, None, None, None\n",
    "\n",
    "dataset_info = [get_dataset_info(dataset.id) for dataset in datasets]\n",
    "dataset_info[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get first row of data data from the datasets-server\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "not enough values to unpack (expected 5, got 4)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[174], line 15\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     13\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m---> 15\u001b[0m dataset_data \u001b[38;5;241m=\u001b[39m [get_dataset_data(dataset\u001b[38;5;241m.\u001b[39mid, config, split, random\u001b[38;5;241m.\u001b[39mchoice(\u001b[38;5;28mlist\u001b[39m(\u001b[38;5;28mrange\u001b[39m(size)))) \u001b[38;5;28;01mfor\u001b[39;00m dataset, (\u001b[38;5;28mid\u001b[39m, config, split, size, features) \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(datasets, dataset_info)]\n\u001b[1;32m     16\u001b[0m dataset_data[\u001b[38;5;241m0\u001b[39m]\n",
      "Cell \u001b[0;32mIn[174], line 15\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     13\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m---> 15\u001b[0m dataset_data \u001b[38;5;241m=\u001b[39m [get_dataset_data(dataset\u001b[38;5;241m.\u001b[39mid, config, split, random\u001b[38;5;241m.\u001b[39mchoice(\u001b[38;5;28mlist\u001b[39m(\u001b[38;5;28mrange\u001b[39m(size)))) \u001b[38;5;28;01mfor\u001b[39;00m dataset, (\u001b[38;5;28mid\u001b[39m, config, split, size, features) \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(datasets, dataset_info)]\n\u001b[1;32m     16\u001b[0m dataset_data[\u001b[38;5;241m0\u001b[39m]\n",
      "\u001b[0;31mValueError\u001b[0m: not enough values to unpack (expected 5, got 4)"
     ]
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "def get_dataset_data(dataset_id, config, split, offset=0, length=1):\n",
    "    if dataset_id and config and split:\n",
    "        url = f\"https://datasets-server.huggingface.co/rows?dataset={dataset_id}&config={config}&split={split}&offset={offset}&length={length}\"\n",
    "        response = requests.get(url)\n",
    "        data = response.json()\n",
    "        if data.get(\"rows\"):\n",
    "            return [row[\"row\"] for row in data[\"rows\"]]\n",
    "        else:\n",
    "            return None\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "dataset_data = [get_dataset_data(dataset.id, config, split, random.choice(list(range(size)))) for dataset, (id, config, split, size, features) in zip(datasets, dataset_info)]\n",
    "dataset_data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Synthesize data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct synthetic SQL Schemas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phi.assistant import Assistant\n",
    "from phi.llm.ollama import Ollama\n",
    "\n",
    "prompt_table_schema = \"\"\"Write a SQL table schema, given a JSON schema and a row of example data. Only return the schema as plain text and nothing else.\n",
    "\n",
    "Schema:\n",
    "{features}\n",
    "\n",
    "Example data:\n",
    "{data}\n",
    "\"\"\"\n",
    "\n",
    "def get_table_schema(features, data):\n",
    "    if features and data:\n",
    "        assistant = Assistant(\n",
    "            llm=Ollama(model=\"llama3.1\"),\n",
    "            show_tool_calls=False,\n",
    "            knowledge_base=knowledge_base\n",
    "        )\n",
    "        return assistant.run(prompt_table_schema.format(features=features, data=data), stream=False)\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "# sql_schemas: list[Iterator[str] | str | BaseModel | None] = [get_table_schema(features[-1], data) for features, data in zip(dataset_info, dataset_data)]\n",
    "table_schema = get_table_schema(features=dataset_info[0][-1], data=dataset_data[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct synthetic natural language queries "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'What is the height of Bronte Tower in feet, and what rank would it hold among the tallest buildings in New York City as of August 2024?'"
      ]
     },
     "execution_count": 172,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nsql_query_prompt = \"\"\"Come up with a reasonable, short and concise natural language query that can be answered using a SQL query on top of the table. Also consider basic aggregates and other operations. Do this using a table schema, a table description, and example data, but consider this is only a single row in the dataset. Only return the natural language query as plain text and nothing else.\n",
    "\n",
    "Description:\n",
    "{description}\n",
    "\n",
    "Schema:\n",
    "{schema}\n",
    "\n",
    "Example data:\n",
    "{data}\n",
    "\"\"\"\n",
    "\n",
    "def get_nsql_query(schema, description, data):\n",
    "    assistant = Assistant(\n",
    "        llm=Ollama(model=\"llama3.1\"),\n",
    "        show_tool_calls=False,\n",
    "        knowledge_base=knowledge_base\n",
    "    )\n",
    "    return assistant.run(nsql_query_prompt.format(schema=schema, description=description, data=data[1]), stream=False)\n",
    "    \n",
    "nsql = get_nsql_query(schema=table_schema, description=card_content[0], data=dataset_data[0])\n",
    "nsql"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct synthetic SQL queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'# FRAMES: Factuality, Retrieval, And reasoning MEasurement Set\\n\\nFRAMES is a comprehensive evaluation dataset designed to test the capabilities of Retrieval-Augmented Generation (RAG) systems across factuality, retrieval accuracy, and reasoning.\\nOur paper with details and experiments is available on arXiv: [https://arxiv.org/abs/2409.12941](https://arxiv.org/abs/2409.12941).\\n\\n\\n## Dataset Overview\\n\\n- 824 challenging multi-hop questions requiring information from 2-15 Wikipedia articles\\n- Questions span diverse topics including history, sports, science, animals, health, etc.\\n- Each question is labeled with reasoning types: numerical, tabular, multiple constraints, temporal, and post-processing\\n- Gold answers and relevant Wikipedia articles provided for each question\\n\\n## Key Features\\n\\n- Tests end-to-end RAG capabilities in a unified framework\\n- Requires integration of information from multiple sources\\n- Incorporates complex reasoning and temporal disambiguation\\n- Designed to be challenging for state-of-the-art language models\\n\\n## Usage\\n\\nThis dataset can be used to:\\n- Evaluate RAG system performance \\n- Benchmark language model factuality and reasoning\\n- Develop and test multi-hop retrieval strategies\\n\\n## Baseline Results\\n\\nWe provide baseline results using state-of-the-art models like Gemini-Pro-1.5-0514:\\n\\n- Naive prompting: 40.8% accuracy\\n- BM25 retrieval (4 docs): 47.4% accuracy  \\n- Oracle retrieval: 72.9% accuracy\\n- Multi-step retrieval & reasoning: 66% accuracy\\n\\n## Citation\\n\\nIf you use this dataset in your research, please cite our paper:\\n\\n```\\n@misc{krishna2024factfetchreasonunified,\\n      title={Fact, Fetch, and Reason: A Unified Evaluation of Retrieval-Augmented Generation}, \\n      author={Satyapriya Krishna and Kalpesh Krishna and Anhad Mohananey and Steven Schwarcz and Adam Stambler and Shyam Upadhyay and Manaal Faruqui},\\n      year={2024},\\n      eprint={2409.12941},\\n      archivePrefix={arXiv},\\n      primaryClass={cs.CL},\\n      url={https://arxiv.org/abs/2409.12941}, \\n}\\n```\\n\\nWe hope FRAMES will be useful for advancing RAG systems and language model capabilities. For more details, please refer to our full paper.'"
      ]
     },
     "execution_count": 166,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "card_content[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"SELECT T1.wikipedia_link_3 \\nFROM table_name AS T1 \\nJOIN table_name AS T2 ON T1.Unnamed_0 = T2.Unnamed_0 \\nWHERE T2.Prompt LIKE '%Dewey Decimal Classification%' AND T2.Answer LIKE '%Jane Eyre%' AND T2.reasoning_types LIKE 'numerical';\""
      ]
     },
     "execution_count": 169,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nsql_to_sql_prompt = \"\"\"Write a valid and fully complete DuckDB SQL query that would be able to provide a good answer to a natural language query, given a dataset scription, table schema and example data. Think step by step and make sure to include all relevant columns and potential statements. Only return the SQL query as plain text and nothing else.\n",
    "\n",
    "Description:\n",
    "{description}\n",
    "\n",
    "Schema:\n",
    "{schema}\n",
    "\n",
    "Natural language query:\n",
    "{nsql}\n",
    "\"\"\"\n",
    "\n",
    "def get_nsql_to_sql(description, schema, data, nsql):\n",
    "    assistant = Assistant(\n",
    "        llm=Ollama(model=\"llama3.1\"),\n",
    "        show_tool_calls=False,\n",
    "        knowledge_base=knowledge_base,\n",
    "        system_prompt=\"You are a SQL expert that understand business scenarios and complexities of natural language question.\"\n",
    "    )\n",
    "    return assistant.run(nsql_to_sql_prompt.format(schema=schema, description=description, data=data[0], nsql=nsql), stream=False)\n",
    "\n",
    "get_nsql_to_sql(schema=table_schema, description=card_content[0], data=dataset_data[0], nsql=nsql)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'row_idx': 0,\n",
       " 'row': {'Unnamed: 0': 0,\n",
       "  'Prompt': \"If my future wife has the same first name as the 15th first lady of the United States' mother and her surname is the same as the second assassinated president's mother's maiden name, what is my future wife's name? \",\n",
       "  'Answer': 'Jane Ballou',\n",
       "  'wikipedia_link_1': 'https://en.wikipedia.org/wiki/President_of_the_United_States',\n",
       "  'wikipedia_link_2': 'https://en.wikipedia.org/wiki/James_Buchanan',\n",
       "  'wikipedia_link_3': 'https://en.wikipedia.org/wiki/Harriet_Lane',\n",
       "  'wikipedia_link_4': 'https://en.wikipedia.org/wiki/List_of_presidents_of_the_United_States_who_died_in_office',\n",
       "  'wikipedia_link_5': 'https://en.wikipedia.org/wiki/James_A._Garfield',\n",
       "  'wikipedia_link_6': None,\n",
       "  'wikipedia_link_7': None,\n",
       "  'wikipedia_link_8': None,\n",
       "  'wikipedia_link_9': None,\n",
       "  'wikipedia_link_10': None,\n",
       "  'wikipedia_link_11+': None,\n",
       "  'reasoning_types': 'Multiple constraints',\n",
       "  'wiki_links': \"['https://en.wikipedia.org/wiki/President_of_the_United_States', 'https://en.wikipedia.org/wiki/James_Buchanan', 'https://en.wikipedia.org/wiki/Harriet_Lane', 'https://en.wikipedia.org/wiki/List_of_presidents_of_the_United_States_who_died_in_office', 'https://en.wikipedia.org/wiki/James_A._Garfield']\"},\n",
       " 'truncated_cells': []}"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validate data Argilla\n",
    "\n",
    "### Create Argilla dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argilla as rg\n",
    "\n",
    "settings = rg.Settings(\n",
    "    fields=[\n",
    "        rg.TextField(name=\"nlq\", title=\"Natural Language Query\"),\n",
    "        rg.TextField(name=\"sql\", title=\"SQL Query\", use_markdown=True),\n",
    "        rg.TextField(name=\"dataset_viewer\", title=\"Dataset viewer used to validate the query.\", use_markdown=True),\n",
    "        rg.TextField(name=\"schema\", title=\"Database schema\", use_markdown=True),\n",
    "        rg.TextField(name=\"description\", title=\"Dataset description\", use_markdown=True),\n",
    "    ],\n",
    "    questions=[\n",
    "        rg.TextQuestion(name=\"nlq\", title=\"Natural Language Query\", description=\"Check if the query makes sense and correct it otherwise.\"),\n",
    "        rg.TextQuestion(name=\"sql\", title=\"SQL Query\", description=\"Check if the SQL query is valid and correct if not use the dataset viewer to iterated and paste the improved SQL query here.\", use_markdown=True)\n",
    "    ]\n",
    ")\n",
    "\n",
    "dataset = rg.Dataset(name=\"text-to-sql\", workspace=\"argilla\", settings=settings)\n",
    "dataset.create()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload data to argilla"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import urllib\n",
    "\n",
    "def get_dataset_viewer(dataset_id, sql_query):\n",
    "    sql_query = urllib.parse.quote(sql_query)\n",
    "    url = f\"https://huggingface.co/datasets/{dataset_id}/embed/viewer?sql_console=true&sql={sql_query}\"\n",
    "    iframe = f\"\"\"\n",
    "    <iframe\n",
    "  src=\"{url}\"\n",
    "  frameborder=\"0\"\n",
    "  width=\"100%\"\n",
    "  height=\"800px\"\n",
    "></iframe>\n",
    "\"\"\"\n",
    "    return iframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "records = [\n",
    "    rg.Record(\n",
    "        fields={\n",
    "            \"schema\": \"\",\n",
    "            \"nlq\": \"\",\n",
    "            \"sql\": \"\",\n",
    "            \"dataset_viewer\": \"\"\n",
    "        },\n",
    "        question={\n",
    "            \"nlq\": \"\",\n",
    "            \"sql\": \"\"\n",
    "        }\n",
    "    ) for schema, \n",
    "]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".env",
   "language": "python",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}