{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d0df5ec3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: gradio in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from -r requirements.txt (line 1)) (5.29.0)\n",
      "Requirement already satisfied: transformers in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from -r requirements.txt (line 2)) (4.51.3)\n",
      "Requirement already satisfied: sentence-transformers in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from -r requirements.txt (line 3)) (4.1.0)\n",
      "Requirement already satisfied: faiss-cpu in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from -r requirements.txt (line 4)) (1.11.0)\n",
      "Requirement already satisfied: datasets in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from -r requirements.txt (line 5)) (3.5.1)\n",
      "Requirement already satisfied: huggingface_hub in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from -r requirements.txt (line 6)) (0.30.2)\n",
      "Requirement already satisfied: aiofiles<25.0,>=22.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (24.1.0)\n",
      "Requirement already satisfied: anyio<5.0,>=3.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (4.8.0)\n",
      "Requirement already satisfied: audioop-lts<1.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.2.1)\n",
      "Requirement already satisfied: fastapi<1.0,>=0.115.2 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.115.12)\n",
      "Requirement already satisfied: ffmpy in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.5.0)\n",
      "Requirement already satisfied: gradio-client==1.10.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (1.10.0)\n",
      "Requirement already satisfied: groovy~=0.1 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.1.2)\n",
      "Requirement already satisfied: httpx>=0.24.1 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.28.1)\n",
      "Requirement already satisfied: jinja2<4.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (3.1.5)\n",
      "Requirement already satisfied: markupsafe<4.0,>=2.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (3.0.2)\n",
      "Requirement already satisfied: numpy<3.0,>=1.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (2.2.3)\n",
      "Requirement already satisfied: orjson~=3.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (3.10.18)\n",
      "Requirement already satisfied: packaging in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (24.2)\n",
      "Requirement already satisfied: pandas<3.0,>=1.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (2.2.3)\n",
      "Requirement already satisfied: pillow<12.0,>=8.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (11.1.0)\n",
      "Requirement already satisfied: pydantic<2.12,>=2.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (2.11.4)\n",
      "Requirement already satisfied: pydub in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.25.1)\n",
      "Requirement already satisfied: python-multipart>=0.0.18 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.0.20)\n",
      "Requirement already satisfied: pyyaml<7.0,>=5.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (6.0.2)\n",
      "Requirement already satisfied: ruff>=0.9.3 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.11.8)\n",
      "Requirement already satisfied: safehttpx<0.2.0,>=0.1.6 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.1.6)\n",
      "Requirement already satisfied: semantic-version~=2.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (2.10.0)\n",
      "Requirement already satisfied: starlette<1.0,>=0.40.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.46.2)\n",
      "Requirement already satisfied: tomlkit<0.14.0,>=0.12.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.13.2)\n",
      "Requirement already satisfied: typer<1.0,>=0.12 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.15.3)\n",
      "Requirement already satisfied: typing-extensions~=4.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (4.12.2)\n",
      "Requirement already satisfied: uvicorn>=0.14.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio->-r requirements.txt (line 1)) (0.34.2)\n",
      "Requirement already satisfied: fsspec in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio-client==1.10.0->gradio->-r requirements.txt (line 1)) (2025.3.0)\n",
      "Requirement already satisfied: websockets<16.0,>=10.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from gradio-client==1.10.0->gradio->-r requirements.txt (line 1)) (15.0.1)\n",
      "Requirement already satisfied: idna>=2.8 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from anyio<5.0,>=3.0->gradio->-r requirements.txt (line 1)) (3.10)\n",
      "Requirement already satisfied: sniffio>=1.1 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from anyio<5.0,>=3.0->gradio->-r requirements.txt (line 1)) (1.3.1)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from pandas<3.0,>=1.0->gradio->-r requirements.txt (line 1)) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from pandas<3.0,>=1.0->gradio->-r requirements.txt (line 1)) (2025.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from pandas<3.0,>=1.0->gradio->-r requirements.txt (line 1)) (2025.1)\n",
      "Requirement already satisfied: annotated-types>=0.6.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from pydantic<2.12,>=2.0->gradio->-r requirements.txt (line 1)) (0.7.0)\n",
      "Requirement already satisfied: pydantic-core==2.33.2 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from pydantic<2.12,>=2.0->gradio->-r requirements.txt (line 1)) (2.33.2)\n",
      "Requirement already satisfied: typing-inspection>=0.4.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from pydantic<2.12,>=2.0->gradio->-r requirements.txt (line 1)) (0.4.0)\n",
      "Requirement already satisfied: click>=8.0.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from typer<1.0,>=0.12->gradio->-r requirements.txt (line 1)) (8.1.8)\n",
      "Requirement already satisfied: shellingham>=1.3.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from typer<1.0,>=0.12->gradio->-r requirements.txt (line 1)) (1.5.4)\n",
      "Requirement already satisfied: rich>=10.11.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from typer<1.0,>=0.12->gradio->-r requirements.txt (line 1)) (14.0.0)\n",
      "Requirement already satisfied: filelock in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from transformers->-r requirements.txt (line 2)) (3.18.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from transformers->-r requirements.txt (line 2)) (2024.11.6)\n",
      "Requirement already satisfied: requests in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from transformers->-r requirements.txt (line 2)) (2.32.3)\n",
      "Requirement already satisfied: tokenizers<0.22,>=0.21 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from transformers->-r requirements.txt (line 2)) (0.21.1)\n",
      "Requirement already satisfied: safetensors>=0.4.3 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from transformers->-r requirements.txt (line 2)) (0.5.3)\n",
      "Requirement already satisfied: tqdm>=4.27 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from transformers->-r requirements.txt (line 2)) (4.67.1)\n",
      "Requirement already satisfied: torch>=1.11.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from sentence-transformers->-r requirements.txt (line 3)) (2.7.0)\n",
      "Requirement already satisfied: scikit-learn in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from sentence-transformers->-r requirements.txt (line 3)) (1.6.1)\n",
      "Requirement already satisfied: scipy in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from sentence-transformers->-r requirements.txt (line 3)) (1.15.2)\n",
      "Requirement already satisfied: pyarrow>=15.0.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from datasets->-r requirements.txt (line 5)) (20.0.0)\n",
      "Requirement already satisfied: dill<0.3.9,>=0.3.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from datasets->-r requirements.txt (line 5)) (0.3.8)\n",
      "Requirement already satisfied: xxhash in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from datasets->-r requirements.txt (line 5)) (3.5.0)\n",
      "Requirement already satisfied: multiprocess<0.70.17 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from datasets->-r requirements.txt (line 5)) (0.70.16)\n",
      "Requirement already satisfied: aiohttp in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from datasets->-r requirements.txt (line 5)) (3.11.18)\n",
      "Requirement already satisfied: hf-xet>=0.1.4 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from huggingface_hub[hf_xet]->-r requirements.txt (line 7)) (1.1.0)\n",
      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from aiohttp->datasets->-r requirements.txt (line 5)) (2.6.1)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from aiohttp->datasets->-r requirements.txt (line 5)) (1.3.2)\n",
      "Requirement already satisfied: attrs>=17.3.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from aiohttp->datasets->-r requirements.txt (line 5)) (25.1.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from aiohttp->datasets->-r requirements.txt (line 5)) (1.6.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from aiohttp->datasets->-r requirements.txt (line 5)) (6.4.3)\n",
      "Requirement already satisfied: propcache>=0.2.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from aiohttp->datasets->-r requirements.txt (line 5)) (0.3.1)\n",
      "Requirement already satisfied: yarl<2.0,>=1.17.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from aiohttp->datasets->-r requirements.txt (line 5)) (1.20.0)\n",
      "Requirement already satisfied: colorama in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from click>=8.0.0->typer<1.0,>=0.12->gradio->-r requirements.txt (line 1)) (0.4.6)\n",
      "Requirement already satisfied: certifi in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from httpx>=0.24.1->gradio->-r requirements.txt (line 1)) (2025.1.31)\n",
      "Requirement already satisfied: httpcore==1.* in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from httpx>=0.24.1->gradio->-r requirements.txt (line 1)) (1.0.7)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from httpcore==1.*->httpx>=0.24.1->gradio->-r requirements.txt (line 1)) (0.14.0)\n",
      "Requirement already satisfied: six>=1.5 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from python-dateutil>=2.8.2->pandas<3.0,>=1.0->gradio->-r requirements.txt (line 1)) (1.17.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from requests->transformers->-r requirements.txt (line 2)) (3.4.1)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from requests->transformers->-r requirements.txt (line 2)) (2.3.0)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from rich>=10.11.0->typer<1.0,>=0.12->gradio->-r requirements.txt (line 1)) (3.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from rich>=10.11.0->typer<1.0,>=0.12->gradio->-r requirements.txt (line 1)) (2.19.1)\n",
      "Requirement already satisfied: mdurl~=0.1 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from markdown-it-py>=2.2.0->rich>=10.11.0->typer<1.0,>=0.12->gradio->-r requirements.txt (line 1)) (0.1.2)\n",
      "Requirement already satisfied: sympy>=1.13.3 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from torch>=1.11.0->sentence-transformers->-r requirements.txt (line 3)) (1.14.0)\n",
      "Requirement already satisfied: networkx in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from torch>=1.11.0->sentence-transformers->-r requirements.txt (line 3)) (3.4.2)\n",
      "Requirement already satisfied: setuptools in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from torch>=1.11.0->sentence-transformers->-r requirements.txt (line 3)) (75.8.2)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from sympy>=1.13.3->torch>=1.11.0->sentence-transformers->-r requirements.txt (line 3)) (1.3.0)\n",
      "Requirement already satisfied: joblib>=1.2.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from scikit-learn->sentence-transformers->-r requirements.txt (line 3)) (1.4.2)\n",
      "Requirement already satisfied: threadpoolctl>=3.1.0 in c:\\users\\victo\\appdata\\local\\programs\\python\\python313\\lib\\site-packages (from scikit-learn->sentence-transformers->-r requirements.txt (line 3)) (3.6.0)\n"
     ]
    }
   ],
   "source": [
    "! pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a50bbe16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RAG configs: ['text-corpus', 'question-answer']\n",
      "Loaded 3200 wiki passages.\n",
      "Loaded 1000 SQuAD examples.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "72faca76fa9f4dd19d6d334cec19c88f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/26 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "09b30d351c304fb9a4cb58ed0fd0fcb6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/26 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 1000 TriviaQA examples.\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset, get_dataset_config_names\n",
    "\n",
    "# 1) See which configs RAG mini-wiki offers\n",
    "print(\"RAG configs:\", get_dataset_config_names(\"rag-datasets/rag-mini-wikipedia\"))\n",
    "\n",
    "# 2) Load the raw wiki passages\n",
    "wiki_ds = load_dataset(\n",
    "    \"rag-datasets/rag-mini-wikipedia\",\n",
    "    \"text-corpus\",\n",
    "    split=\"passages\"\n",
    ")\n",
    "passages = wiki_ds[\"passage\"]\n",
    "print(f\"Loaded {len(passages)} wiki passages.\")\n",
    "\n",
    "# 3) Load a small SQuAD v2 slice (first 1 000 examples for speed)\n",
    "squad = load_dataset(\"rajpurkar/squad_v2\", split=\"train[:1000]\")\n",
    "print(f\"Loaded {len(squad)} SQuAD examples.\")\n",
    "\n",
    "# 4) Load TriviaQA small\n",
    "trivia = load_dataset(\n",
    "    \"mandarjoshi/trivia_qa\",\n",
    "    \"rc\",\n",
    "    split=\"validation[:1000]\"\n",
    ")\n",
    "print(f\"Loaded {len(trivia)} TriviaQA examples.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e9307bab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2ee1f7b78acf419eb853d10599b86438",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexed 3200 vectors of size 384.\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "import faiss\n",
    "\n",
    "# Embedder\n",
    "embedder = SentenceTransformer(\"all-MiniLM-L6-v2\")\n",
    "embeddings = embedder.encode(passages, show_progress_bar=True, convert_to_numpy=True)\n",
    "\n",
    "# Build index\n",
    "dim = embeddings.shape[1]\n",
    "index = faiss.IndexFlatL2(dim)\n",
    "index.add(embeddings)\n",
    "print(f\"Indexed {index.ntotal} vectors of size {dim}.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "745eb30e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Device set to use cpu\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a neural network\n"
     ]
    }
   ],
   "source": [
    "# ==== 4. Load & Test the LLM ====\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline\n",
    "\n",
    "# 4.1) Use the correct Flan-T5 repo\n",
    "MODEL_NAME = \"google/flan-t5-base\"\n",
    "\n",
    "# 4.2) Load tokenizer & model (requires that you ran `huggingface-cli login`)\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)\n",
    "\n",
    "# 4.3) Create a generation pipeline\n",
    "qa_pipeline = pipeline(\n",
    "    \"text2text-generation\",\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    device=-1            # or 0 if you have a GPU\n",
    ")\n",
    "\n",
    "# 4.4) Quick sanity check\n",
    "print(qa_pipeline(\"Question: What is retrieval-augmented generation? Answer:\", max_length=50)[0][\"generated_text\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "70efc796",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer: Isaac Newton\n"
     ]
    }
   ],
   "source": [
    "# ==== 5. Retrieval + Generation Function ====\n",
    "def retrieve_and_answer(question: str, k: int = 5) -> str:\n",
    "    # 1) Embed the question\n",
    "    q_emb = embedder.encode([question], convert_to_numpy=True)\n",
    "    # 2) Retrieve top-k passages\n",
    "    distances, indices = index.search(q_emb, k)\n",
    "    # 3) Build a little context block\n",
    "    contexts = \"\\n\".join(\n",
    "        f\"Context {i+1}: {passages[idx][:200]}...\"\n",
    "        for i, idx in enumerate(indices[0])\n",
    "    )\n",
    "    # 4) Assemble prompt\n",
    "    prompt = (\n",
    "        \"You are a helpful QA assistant. \"\n",
    "        \"Use ONLY the following contexts to answer the question. \"\n",
    "        \"If the answer is not contained in the contexts, respond with \"\n",
    "        \"'Sorry, I don't know.'\\n\\n\"\n",
    "        f\"{contexts}\\n\"\n",
    "        f\"Question: {question}\\nAnswer:\"\n",
    "    )\n",
    "\n",
    "    # 5) Generate and return\n",
    "    out = qa_pipeline(prompt, max_length=200, do_sample=False)\n",
    "    return out[0][\"generated_text\"]\n",
    "\n",
    "\n",
    "# 5.1) Test it end-to-end\n",
    "sample_q = \"Who conceptualized the theory of relativity?\"\n",
    "print(\"Answer:\", retrieve_and_answer(sample_q))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a1476c3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "\n",
    "def chat_fn(user_message, history=None):\n",
    "    # Initialize history if needed\n",
    "    history = history or []\n",
    "    # 1) Record the user turn\n",
    "    history.append({\"role\": \"user\", \"content\": user_message})\n",
    "    # 2) Retrieve & generate\n",
    "    answer = retrieve_and_answer(user_message, k=5)\n",
    "    # 3) Record the assistant turn\n",
    "    history.append({\"role\": \"assistant\", \"content\": answer})\n",
    "    # 4) Return the full message list\n",
    "    return history\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58d445a5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "fc557ef9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Running on local URL:  http://127.0.0.1:7865\n",
      "\n",
      "Could not create share link. Missing file: C:\\Users\\victo\\.cache\\huggingface\\gradio\\frpc\\frpc_windows_amd64_v0.3. \n",
      "\n",
      "Please check your internet connection. This can happen if your antivirus software blocks the download of this file. You can install manually by following these steps: \n",
      "\n",
      "1. Download this file: https://cdn-media.huggingface.co/frpc-gradio-0.3/frpc_windows_amd64.exe\n",
      "2. Rename the downloaded file to: frpc_windows_amd64_v0.3\n",
      "3. Move the file to this location: C:\\Users\\victo\\.cache\\huggingface\\gradio\\frpc\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7865/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Launch with the new format\n",
    "gr.ChatInterface(\n",
    "    fn=chat_fn,\n",
    "    type=\"messages\",            # ← switch to messages!\n",
    "    title=\"🔍 RAG QA Demo\",\n",
    "    description=\"I only answer from retrieved contexts; otherwise I'll say 'Sorry, I don't know.'\"\n",
    ").launch(share=True)            # ← add share=True if you want a public link"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ec42db4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}