refactoring

.env

@@ -0,0 +1,20 @@
+export TABLE_NAME=minilm_table
+export EMB_MODEL=sentence-transformers/all-MiniLM-L6-v2
+export TOP_K=5
+export HF_MODEL=mistralai/Mistral-7B-Instruct-v0.2
+export OPENAI_MODEL=gpt-4-turbo-preview
+
+#### SECRETS ####
+export OPENAI_API_KEY=
+export HF_TOKEN=
+################
+
+export TEMPERATURE=
+export MAX_NEW_TOKENS=
+export TOP_P=
+export REP_PENALTY=
+export DO_SAMPLE=
+export FREQ_PENALTY=
+export VECTOR_COLUMN=
+export TEXT_COLUMN=
+export BATCH_SIZE=

README.md

@@ -1,12 +1,7 @@
----
-title: Test
-emoji: 🏃
-colorFrom: red
-colorTo: indigo
-sdk: gradio
-sdk_version: 4.16.0
-app_file: app.py
-pinned: false
----
+## A template for a RAG system with Gradio UI
+Deliberately stripped down for the sake of an exercise.
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+For https://ydataschool.com/
+
+## Setup
+Clone the repo and follow **rag-homework.ipynb**

gradio_app/app.py

@@ -1,11 +1,7 @@
 """
 Credit to Derek Thomas, derek@huggingface.co
 """
-
-import subprocess
-
-subprocess.run(["pip", "install", "--upgrade", "transformers[torch,sentencepiece]==4.34.1"])
-
+import os
 import logging
 from pathlib import Path
 from time import perf_counter
@@ -14,10 +10,10 @@ import gradio as gr
 from jinja2 import Environment, FileSystemLoader
 
 from backend.query_llm import generate_hf, generate_openai
-from backend.semantic_search import table, retriever
+from backend.semantic_search import retrieve
 
-VECTOR_COLUMN_NAME = ""
-TEXT_COLUMN_NAME = ""
+
+TOP_K = int(os.getenv("TOP_K", 4))
 
 proj_dir = Path(__file__).parent
 # Setting up the logging
@@ -31,11 +27,6 @@ env = Environment(loader=FileSystemLoader(proj_dir / 'templates'))
 template = env.get_template('template.j2')
 template_html = env.get_template('template_html.j2')
 
-# Examples
-examples = ['What is the capital of China?',
-            'Why is the sky blue?',
-            'Who won the mens world cup in 2014?', ]
-
 
 def add_text(history, text):
     history = [] if history is None else history
@@ -44,23 +35,19 @@ def add_text(history, text):
 
 
 def bot(history, api_kind):
-    top_k_rank = 4
     query = history[-1][0]
 
     if not query:
-         gr.Warning("Please submit a non-empty string as a prompt")
-         raise ValueError("Empty string was submitted")
+        raise gr.Warning("Please submit a non-empty string as a prompt")
 
-    logger.warning('Retrieving documents...')
+    logger.info('Retrieving documents...')
     # Retrieve documents relevant to query
     document_start = perf_counter()
 
-    query_vec = retriever.encode(query)
-    documents = table.search(query_vec, vector_column_name=VECTOR_COLUMN_NAME).limit(top_k_rank).to_list()
-    documents = [doc[TEXT_COLUMN_NAME] for doc in documents]
+    documents = retrieve(query, TOP_K)
 
     document_time = perf_counter() - document_start
-    logger.warning(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
+    logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
 
     # Create Prompt
     prompt = template.render(documents=documents, query=query)
@@ -70,12 +57,8 @@ def bot(history, api_kind):
          generate_fn = generate_hf
     elif api_kind == "OpenAI":
          generate_fn = generate_openai
-    elif api_kind is None:
-         gr.Warning("API name was not provided")
-         raise ValueError("API name was not provided")
     else:
-         gr.Warning(f"API {api_kind} is not supported")
-         raise ValueError(f"API {api_kind} is not supported")
+         raise gr.Error(f"API {api_kind} is not supported")
 
     history[-1][1] = ""
     for character in generate_fn(prompt, history[:-1]):
@@ -120,8 +103,5 @@ with gr.Blocks() as demo:
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
 
-    # Examples
-    gr.Examples(examples, txt)
-
 demo.queue()
 demo.launch(debug=True)

gradio_app/backend/query_llm.py

@@ -1,25 +1,37 @@
 import openai
 import gradio as gr
+import os
 
-from os import getenv
 from typing import Any, Dict, Generator, List
 
 from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
 
-temperature = 0.9
-top_p = 0.6
-repetition_penalty = 1.2
+OPENAI_KEY = os.getenv("OPENAI_API_KEY")
+HF_TOKEN = os.getenv("HF_TOKEN")
+TOKENIZER = AutoTokenizer.from_pretrained(os.getenv("HF_MODEL"))
 
-OPENAI_KEY = getenv("OPENAI_API_KEY")
-HF_TOKEN = getenv("HUGGING_FACE_HUB_TOKEN")
+HF_CLIENT = InferenceClient(
+    os.getenv("HF_MODEL"),
+    token=HF_TOKEN
+)
+OAI_CLIENT = openai.Client(api_key=OPENAI_KEY)
 
-hf_client = InferenceClient(
-        "mistralai/Mistral-7B-Instruct-v0.1",
-        token=HF_TOKEN
-        )
+HF_GENERATE_KWARGS = {
+    'temperature': max(float(os.getenv("TEMPERATURE", 0.9)), 1e-2),
+    'max_new_tokens': int(os.getenv("MAX_NEW_TOKENS", 256)),
+    'top_p': float(os.getenv("TOP_P", 0.6)),
+    'repetition_penalty': float(os.getenv("REP_PENALTY", 1.2)),
+    'do_sample': bool(os.getenv("DO_SAMPLE", True))
+}
+
+OAI_GENERATE_KWARGS = {
+    'temperature': max(float(os.getenv("TEMPERATURE", 0.9)), 1e-2),
+    'max_tokens': int(os.getenv("MAX_NEW_TOKENS", 256)),
+    'top_p': float(os.getenv("TOP_P", 0.6)),
+    'frequency_penalty': max(-2, min(float(os.getenv("FREQ_PENALTY", 0)), 2))
+}
 
 
 def format_prompt(message: str, api_kind: str):
@@ -28,7 +40,7 @@ def format_prompt(message: str, api_kind: str):
 
     Args:
         message (str): The user message to be formatted.
-
+        api_kind (str): LLM API provider.
     Returns:
         str: Formatted message after applying the chat template.
     """
@@ -39,46 +51,34 @@ def format_prompt(message: str, api_kind: str):
     if api_kind == "openai":
         return messages
     elif api_kind == "hf":
-        return tokenizer.apply_chat_template(messages, tokenize=False)
+        return TOKENIZER.apply_chat_template(messages, tokenize=False)
     elif api_kind:
         raise ValueError("API is not supported")
 
 
-def generate_hf(prompt: str, history: str, temperature: float = 0.9, max_new_tokens: int = 256,
-             top_p: float = 0.95, repetition_penalty: float = 1.0) -> Generator[str, None, str]:
+def generate_hf(prompt: str, history: str) -> Generator[str, None, str]:
     """
     Generate a sequence of tokens based on a given prompt and history using Mistral client.
 
     Args:
-        prompt (str): The initial prompt for the text generation.
+        prompt (str): The prompt for the text generation.
         history (str): Context or history for the text generation.
-        temperature (float, optional): The softmax temperature for sampling. Defaults to 0.9.
-        max_new_tokens (int, optional): Maximum number of tokens to be generated. Defaults to 256.
-        top_p (float, optional): Nucleus sampling probability. Defaults to 0.95.
-        repetition_penalty (float, optional): Penalty for repeated tokens. Defaults to 1.0.
-
     Returns:
         Generator[str, None, str]: A generator yielding chunks of generated text.
                                    Returns a final string if an error occurs.
     """
 
-    temperature = max(float(temperature), 1e-2)  # Ensure temperature isn't too low
-    top_p = float(top_p)
-
-    generate_kwargs = {
-        'temperature': temperature,
-        'max_new_tokens': max_new_tokens,
-        'top_p': top_p,
-        'repetition_penalty': repetition_penalty,
-        'do_sample': True,
-        'seed': 42,
-        }
-    
     formatted_prompt = format_prompt(prompt, "hf")
+    formatted_prompt = formatted_prompt.encode("utf-8").decode("utf-8")
 
     try:
-        stream = hf_client.text_generation(formatted_prompt, **generate_kwargs,
-                                            stream=True, details=True, return_full_text=False)
+        stream = HF_CLIENT.text_generation(
+            formatted_prompt,
+            **HF_GENERATE_KWARGS,
+            stream=True,
+            details=True,
+            return_full_text=False
+        )
         output = ""
         for response in stream:
             output += response.token.text
@@ -86,69 +86,43 @@ def generate_hf(prompt: str, history: str, temperature: float = 0.9, max_new_tok
 
     except Exception as e:
         if "Too Many Requests" in str(e):
-            print("ERROR: Too many requests on Mistral client")
-            gr.Warning("Unfortunately Mistral is unable to process")
-            return "Unfortunately, I am not able to process your request now."
+            raise gr.Error(f"Too many requests: {str(e)}")
         elif "Authorization header is invalid" in str(e):
-            print("Authetification error:", str(e))
-            gr.Warning("Authentication error: HF token was either not provided or incorrect")
-            return "Authentication error"
+            raise gr.Error("Authentication error: HF token was either not provided or incorrect")
         else:
-            print("Unhandled Exception:", str(e))
-            gr.Warning("Unfortunately Mistral is unable to process")
-            return "I do not know what happened, but I couldn't understand you."
+            raise gr.Error(f"Unhandled Exception: {str(e)}")
 
 
-def generate_openai(prompt: str, history: str, temperature: float = 0.9, max_new_tokens: int = 256,
-             top_p: float = 0.95, repetition_penalty: float = 1.0) -> Generator[str, None, str]:
+def generate_openai(prompt: str, history: str) -> Generator[str, None, str]:
     """
     Generate a sequence of tokens based on a given prompt and history using Mistral client.
 
     Args:
         prompt (str): The initial prompt for the text generation.
         history (str): Context or history for the text generation.
-        temperature (float, optional): The softmax temperature for sampling. Defaults to 0.9.
-        max_new_tokens (int, optional): Maximum number of tokens to be generated. Defaults to 256.
-        top_p (float, optional): Nucleus sampling probability. Defaults to 0.95.
-        repetition_penalty (float, optional): Penalty for repeated tokens. Defaults to 1.0.
-
     Returns:
         Generator[str, None, str]: A generator yielding chunks of generated text.
                                    Returns a final string if an error occurs.
     """
-
-    temperature = max(float(temperature), 1e-2)  # Ensure temperature isn't too low
-    top_p = float(top_p)
-    
-    generate_kwargs = {
-        'temperature': temperature,
-        'max_tokens': max_new_tokens,
-        'top_p': top_p,
-        'frequency_penalty': max(-2., min(repetition_penalty, 2.)),
-        }
-
     formatted_prompt = format_prompt(prompt, "openai")
 
     try:
-        stream = openai.ChatCompletion.create(model="gpt-3.5-turbo-0301",
-                                                messages=formatted_prompt, 
-                                                **generate_kwargs, 
-                                                stream=True)
+        stream = OAI_CLIENT.chat.completions.create(
+            model=os.getenv("OPENAI_MODEL"),
+            messages=formatted_prompt,
+            **OAI_GENERATE_KWARGS,
+            stream=True
+        )
         output = ""
         for chunk in stream:
-            output += chunk.choices[0].delta.get("content", "")
-            yield output
+            if chunk.choices[0].delta.content:
+                output += chunk.choices[0].delta.content
+                yield output
 
     except Exception as e:
         if "Too Many Requests" in str(e):
-            print("ERROR: Too many requests on OpenAI client")
-            gr.Warning("Unfortunately OpenAI is unable to process")
-            return "Unfortunately, I am not able to process your request now."
+            raise gr.Error("ERROR: Too many requests on OpenAI client")
         elif "You didn't provide an API key" in str(e):
-            print("Authetification error:", str(e))
-            gr.Warning("Authentication error: OpenAI key was either not provided or incorrect")
-            return "Authentication error"
+            raise gr.Error("Authentication error: OpenAI key was either not provided or incorrect")
         else:
-            print("Unhandled Exception:", str(e))
-            gr.Warning("Unfortunately OpenAI is unable to process")
-            return "I do not know what happened, but I couldn't understand you."
+            raise gr.Error(f"Unhandled Exception: {str(e)}")

gradio_app/backend/semantic_search.py

@@ -1,18 +1,26 @@
-import logging
 import lancedb
 import os
-from pathlib import Path
+import gradio as gr
 from sentence_transformers import SentenceTransformer
 
-EMB_MODEL_NAME = ""
-DB_TABLE_NAME = ""
 
-# Setting up the logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-retriever = SentenceTransformer(EMB_MODEL_NAME)
+db = lancedb.connect(".lancedb")
 
-# db
-db_uri = os.path.join(Path(__file__).parents[1], ".lancedb")
-db = lancedb.connect(db_uri)
-table = db.open_table(DB_TABLE_NAME)
+TABLE = db.open_table(os.getenv("TABLE_NAME"))
+VECTOR_COLUMN = os.getenv("VECTOR_COLUMN", "vector")
+TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
+BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
+
+retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
+
+
+def retrieve(query, k):
+    query_vec = retriever.encode(query)
+    try:
+        documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
+        documents = [doc[TEXT_COLUMN] for doc in documents]
+
+        return documents
+
+    except Exception as e:
+        raise gr.Error(str(e))

gradio_app/requirements.txt

@@ -1,9 +1,7 @@
-# transformers[torch,sentencepiece]==4.34.1
-wikiextractor==3.0.6
-sentence-transformers>2.2.0
-ipywidgets==8.1.1
+lancedb==0.5.3
+openai==1.11.1
+sentence-transformers==2.3.1
 tqdm==4.66.1
-aiohttp==3.8.6
-huggingface-hub==0.17.3
-lancedb==0.3.1
-openai==0.28
+torch==2.1.1
+transformers==4.37.2
+

prep_scripts/lancedb_setup.py

@@ -1,3 +1,4 @@
+import argparse
 import lancedb
 import torch
 import pyarrow as pa
@@ -5,61 +6,91 @@ import pandas as pd
 from pathlib import Path
 import tqdm
 import numpy as np
+import logging
 
+from transformers import AutoConfig
 from sentence_transformers import SentenceTransformer
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-EMB_MODEL_NAME = ""
-DB_TABLE_NAME = ""
-VECTOR_COLUMN_NAME = ""
-TEXT_COLUMN_NAME = ""
-INPUT_DIR = "<chunked docs directory>"
-db = lancedb.connect(".lancedb") # db location
-batch_size = 32
-
-model = SentenceTransformer(EMB_MODEL_NAME)
-model.eval()
-
-if torch.backends.mps.is_available():
-    device = "mps"
-elif torch.cuda.is_available():
-    device = "cuda"
-else:
-    device = "cpu"
-
-schema = pa.schema(
-  [
-      pa.field(VECTOR_COLUMN_NAME, pa.list_(pa.float32(), 768)),
-      pa.field(TEXT_COLUMN_NAME, pa.string())
-  ])
-tbl = db.create_table(DB_TABLE_NAME, schema=schema, mode="overwrite")
-
-input_dir = Path(INPUT_DIR)
-files = list(input_dir.rglob("*"))
-
-sentences = []
-for file in files:
-    with open(file) as f:
-        sentences.append(f.read())
-
-for i in tqdm.tqdm(range(0, int(np.ceil(len(sentences) / batch_size)))):
-    try:
-        batch = [sent for sent in sentences[i * batch_size:(i + 1) * batch_size] if len(sent) > 0]
-        encoded = model.encode(batch, normalize_embeddings=True, device=device)
-        encoded = [list(vec) for vec in encoded]
-
-        df = pd.DataFrame({
-            VECTOR_COLUMN_NAME: encoded,
-            TEXT_COLUMN_NAME: batch
-        })
-
-        tbl.add(df)
-    except:
-        print(f"batch {i} was skipped")
-
-'''
-create ivf-pd index https://lancedb.github.io/lancedb/ann_indexes/
-with the size of the transformer docs, index is not really needed
-but we'll do it for demonstrational purposes
-'''
-tbl.create_index(num_partitions=256, num_sub_vectors=96, vector_column_name=VECTOR_COLUMN_NAME)
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--emb-model", help="embedding model name on HF hub", type=str)
+    parser.add_argument("--table", help="table name in DB", type=str)
+    parser.add_argument("--input-dir", help="input directory with documents to ingest", type=str)
+    parser.add_argument("--vec-column", help="vector column name in the table", type=str, default="vector")
+    parser.add_argument("--text-column", help="text column name in the table", type=str, default="text")
+    parser.add_argument("--db-loc", help="database location", type=str,
+                        default=str(Path().resolve() / "gradio_app" / ".lancedb"))
+    parser.add_argument("--batch-size", help="batch size for embedding model", type=int, default=32)
+    parser.add_argument("--num-partitions", help="number of partitions for index", type=int, default=256)
+    parser.add_argument("--num-sub-vectors", help="number of sub-vectors for index", type=int, default=96)
+
+    args = parser.parse_args()
+
+    emb_config = AutoConfig.from_pretrained(args.emb_model)
+    emb_dimension = emb_config.hidden_size
+
+    assert emb_dimension % args.num_sub_vectors == 0, \
+        "Embedding size must be divisible by the num of sub vectors"
+
+    model = SentenceTransformer(args.emb_model)
+    model.eval()
+
+    if torch.backends.mps.is_available():
+        device = "mps"
+    elif torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+    logger.info(f"using {str(device)} device")
+
+    db = lancedb.connect(args.db_loc)
+
+    schema = pa.schema(
+      [
+          pa.field(args.vec_column, pa.list_(pa.float32(), emb_dimension)),
+          pa.field(args.text_column, pa.string())
+      ]
+    )
+    tbl = db.create_table(args.table, schema=schema, mode="overwrite")
+
+    input_dir = Path(args.input_dir)
+    files = list(input_dir.rglob("*"))
+
+    sentences = []
+    for file in files:
+        with open(file) as f:
+            sentences.append(f.read())
+
+    for i in tqdm.tqdm(range(0, int(np.ceil(len(sentences) / args.batch_size)))):
+        try:
+            batch = [sent for sent in sentences[i * args.batch_size:(i + 1) * args.batch_size] if len(sent) > 0]
+            encoded = model.encode(batch, normalize_embeddings=True, device=device)
+            encoded = [list(vec) for vec in encoded]
+
+            df = pd.DataFrame({
+                args.vec_column: encoded,
+                args.text_column: batch
+            })
+
+            tbl.add(df)
+        except:
+            logger.info(f"batch {i} was skipped")
+
+    '''
+    create ivf-pd index https://lancedb.github.io/lancedb/ann_indexes/
+    with the size of the transformer docs, index is not really needed
+    but we'll do it for demonstrational purposes
+    '''
+    tbl.create_index(
+        num_partitions=args.num_partitions,
+        num_sub_vectors=args.num_sub_vectors,
+        vector_column_name=args.vec_column
+    )
+
+
+if __name__ == "__main__":
+    main()

prep_scripts/markdown_to_text.py

@@ -1,12 +1,11 @@
-from bs4 import BeautifulSoup
-from markdown import markdown
+import argparse
 import os
 import re
-from pathlib import Path
 
-
-DIR_TO_SCRAPE = "transformers/docs/source/en/"
-OUTPUT_DIR = str(Path().resolve() / "docs_dump")
+from tqdm import tqdm
+from bs4 import BeautifulSoup
+from markdown import markdown
+from pathlib import Path
 
 
 def markdown_to_text(markdown_string):
@@ -31,19 +30,33 @@ def markdown_to_text(markdown_string):
     return text
 
 
-dir_to_scrape = Path(DIR_TO_SCRAPE)
-files = list(dir_to_scrape.rglob("*"))
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input-dir", help="input directory with markdown", type=str,
+                        default="transformers/docs/source/en/")
+    parser.add_argument("--output-dir", help="output directory to store raw texts", type=str,
+                        default="docs")
+
+    args = parser.parse_args()
+    input_dir = Path(args.input_dir)
+    output_dir = Path(args.output_dir)
+
+    assert os.path.isdir(input_dir), "Input directory doesn't exist"
+
+    files = input_dir.rglob("*")
+    os.makedirs(output_dir, exist_ok=True)
 
-os.makedirs(OUTPUT_DIR, exist_ok=True)
+    for file in tqdm(files):
+        parent = file.parent.stem if file.parent.stem != input_dir.stem else ""
+        if file.is_file():
+            with open(file) as f:
+                md = f.read()
 
-for file in files:
-    parent = file.parent.stem if file.parent.stem != dir_to_scrape.stem else ""
-    if file.is_file():
-        with open(file) as f:
-            md = f.read()
+            text = markdown_to_text(md)
 
-        text = markdown_to_text(md)
+            with open(output_dir / f"{parent}_{file.stem}.txt", "w") as f:
+                f.write(text)
 
-        with open(os.path.join(OUTPUT_DIR, f"{parent}_{file.stem}.txt"), "w") as f:
-            f.write(text)
 
+if __name__ == "__main__":
+    main()

rag-homework.ipynb

@@ -0,0 +1,270 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "844fe3af-9cf1-4c66-aa78-b88a3429acc6",
+   "metadata": {},
+   "source": [
+    "### 0. Setup\n",
+    "- Clone https://github.com/plaggy/rag-gradio-sample-project and set up an environment with gradio_app/requirements.txt.\n",
+    "\n",
+    "A convenient way to work through the project is to test locally and keep committing the changes to the [HF Spaces](https://huggingface.co/spaces) repo. A space gets automatically rebuilt after each commit and you get a new version of your application up and running.\n",
+    "\n",
+    "- Create a new space with Gradio SDK. You'll get an almost empty repo, the only thing you'll need from it is README.md which has a config letting a space builder know that it's a Gradio app. Reset a remote upstream of your local rag-gradio-sample-project clone to be your freshly created Spaces repository.\n",
+    "\n",
+    "The easiest way to set your space up is to set up the gradio_app folder as a git repo, set remote origin to your space repo and checkout the remote README:\n",
+    "\n",
+    "```\n",
+    "cd gradio_app\n",
+    "git init\n",
+    "git remote add origin <your spaces repo url>\n",
+    "git fetch\n",
+    "git checkout origin/main README.md\n",
+    "```\n",
+    "\n",
+    "The space is not working yet. You'll get the first working version after the Step 3.\n",
+    "\n",
+    "- Clone https://github.com/huggingface/transformers to a local machine and run prep_scripts/markdown_to_text.py script to extract raw text from transformers/docs/source/en/. This will be your knowledge base, you don't need it to be a part of your repository\n",
+    "\n",
+    "Run the command as follows (pass arguments that work for you)\n",
+    "```\n",
+    "python prep_scripts/markdown_to_text.py --input-dir transformers/docs/source/en/ --output-dir docs\n",
+    "```\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "762e9fde-c1f4-464c-b12b-dca602fac5ba",
+   "metadata": {},
+   "source": [
+    "**By design, you'll be running your experiments in a [Gradio space](https://huggingface.co/docs/hub/en/spaces-sdks-gradio). Apart from deliverables for each step you'll need to provide a link to a functioning RAG space in it final state!**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "6c813d03-33a7-4ce1-836f-11afc541f291",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Add the link to the space you've just created here:"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c970d0a4-fee8-48ac-9377-4a6def7712b2",
+   "metadata": {},
+   "source": [
+    "### Step 1: Chunk Your Data\n",
+    "\n",
+    "To efficiently pull up documents relevant to a query from a knowledge base documents are embedded and stored as vectors. Documents in your knowledge base are not expected to fit into the context length of an embedding model (most have 512 token limit). Hence chunking your documents into smaller pieces is required. Take a deeper dive into why chunking is imoprtant and what are the options [here](https://www.pinecone.io/learn/chunking-strategies/).\n",
+    "\n",
+    "Your task is to implement and compare two chunking strategies: fixed-sized chunking and content-aware chunking. For content-aware you could split by sentences, paragraphs or in some other way that makes sence.\n",
+    "\n",
+    "The deliverables are:\n",
+    "\n",
+    "1. The code for chunk splitting\n",
+    "2. The illustrative examples of how retrived chunks differ depending on the chunking strategy.\n",
+    "3. The analysis of pros and cons of each strategy\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f7bad8c8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# chunk splitting"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5e5ebaad-8d42-430c-b00b-18198cdb9ce8",
+   "metadata": {},
+   "source": [
+    "### Step 2: Ingest chunks into a database and create an index\n",
+    "\n",
+    "Chunks need to be vectorized and made accessible to an LLM to enable semantic search with embedding models. A current industry standard is to use a vector database to store and retrieve texts both conveniently and efficiently. There are many products out there, we'll be using [LanceDB](https://lancedb.github.io/lancedb/). LanceDB is a young product, one way it stands out is that it's embedded - it's designed not to be a standalone service but rather a part of an application, more on this [here](https://lancedb.github.io/lancedb/basic/).\n",
+    "\n",
+    "Your task is to vectorize and ingest chunked documents into the database. Use prep_scrips/lancedb_setup.py to vectorize chunks and store vector representations along with raw text in a Lancedb instance. The script also creates an index for fast ANN retrieval (not really needed for this exercise but necessary at scale). Try different embedding models and see how results differ. The options are:\n",
+    "\n",
+    "- `sentence-transformers/all-MiniLM-L6-v2`: a light model, produces vectors of length 384\n",
+    "- `BAAI/bge-large-en-v1.5`: a much heavier model, embedding vector length is 1024\n",
+    "\n",
+    "Feel free to explore other embedding models and justify your choice.\n",
+    "For different embedding model create different tables in the database so you can easily switch between them and compare.\n",
+    "\n",
+    "Run the embedding+ingestion script as follows, make sure to look into the script and see all the arguments. Note that the number of sub-vectors for indexing must be a divisor of the model embedding size.\n",
+    "\n",
+    "```\n",
+    "python prep_scrips/lancedb_setup.py --emb-model <model name> --table <db table name> --input-dir <folder with docs> --num-sub-vectors <a number which is a divisor of the embedding dim>\n",
+    "```\n",
+    "\n",
+    "Before committing to your space set up environment variables on the settings tab of your space, use `.env` as a ference list of all the things you can customize. Make sure to add HF_TOKEN and OPENAI_API_KEY as secrets.\n",
+    "Not all the parameters are required to set via environment variables, most have default values.\n",
+    "\n",
+    "!!The database is expected to be in the `gradio_app` folder under `.lancedb`, make sure to move it there if was initialized elsewhere. It can be parametrized but it's unnecessary here.\n",
+    "\n",
+    "To commit large files to Github use `git lfs`:\n",
+    "```\n",
+    "git lfs install\n",
+    "git lfs track \"*.lance\"\n",
+    "git lfs track \"*.idx\"\n",
+    "git add .gitattributes\n",
+    "```\n",
+    "Then proceed as usual.\n",
+    "\n",
+    "For experimenting you can easily switch between embedding models/tables by changing the values of the corresponding env variables in your space (EMB_MODEL, TABLE_NAME). Overall, every time you change the value of an environment variable a space gets automatically rebuilt.\n",
+    "\n",
+    "The deliverables are:\n",
+    "1. The illustration of how retrieved documents differ depending on the embedding model\n",
+    "2. The analysis of how the embedding time differs between models\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f7db282e-e03c-41de-9c03-54abf455481f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Embed documents and ingest into the database "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c929fe99-7f3f-49cf-af9f-07bac9668c2f",
+   "metadata": {},
+   "source": [
+    "### Step 3: DB retrieval\n",
+    "\n",
+    "To augment an LLM with a knowledge base the context provided in a prompt is enriched with relevant pieces from a DB. An incoming query is vectorized with an embedding model and top-k relevant pieces are retrieved based on similarity to a query vector. Your task is to implement database retrieval and adjust the prompt to include the documents from the DB.\n",
+    "\n",
+    "The deliverables are:\n",
+    "\n",
+    "1. The new code that enables RAG. Paste the code added and some surrounding context; indicate which file it comes from.\n",
+    "2. A new prompt template\n",
+    "3. Examples of how the output changes\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5bc1399d-1c3b-48b7-a9fb-ed48a0dfaa6a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# DB retrieval implementation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7d818b4f-ba5a-4c81-b6d7-f3474c398d9c",
+   "metadata": {},
+   "source": [
+    "### Step 4: Add a reranker\n",
+    "\n",
+    "A reranker is a second-level model which produces similarity scores for pairs of (input query + retrieved document). Cross-encoders are conventionally used for reranking, their architecture is slightly different from retrieval models (more on it [here]). Cross-encoders are much more costly to run, therefore a retrieval model is used to get a few (dozens) highest-scoring items, and a reranker picks the best among these. The overall pipeline is similar to the recommender system indudustry standard: a light model retrieves top-n, a precise and heavy model reranks n to get top k, n is orders of magnitude larger than k.\n",
+    "\n",
+    "Cross-encoders are optional because of the overhead their usage implies. Your task is to implement a reranker using a cross-encoder and assess pros and cons of having it. Do not forget that the process of pulling the most relevant documents becomes two-staged: retrieve a larger number of items first, than rerank and keep the best top-k for context.\n",
+    "\n",
+    "The models fit for the task:\n",
+    "1. BAAI/bge-reranker-large\n",
+    "2. cross-encoder/ms-marco-MiniLM-L-6-v2\n",
+    "\n",
+    "As usual, feel free to pick another model and provide some description to it.\n",
+    "\n",
+    "The deliverables are:\n",
+    "\n",
+    "1. The code that enables a reranker.\n",
+    "2. A comparison of how the prompt and the model output change after adding a reranker\n",
+    "3. The analysis of pros and cons. The evaluation aspects should include the relevance of the top-k documents, the response time.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ee1b0160-0ba0-4b5f-81c4-ef3ea76850e5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Implement code for selecting the final documents using a cross-encoder"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f5816c54-a290-4cb0-b7db-3b8901998cb0",
+   "metadata": {},
+   "source": [
+    "### Step 5: Try a different LLM\n",
+    "\n",
+    "The suggested Mistral-7b is a great but small model for an LLM. A larger model can be applied to a wider range of problems and do more complex reasoning. Within the scope of this project a larger model may not be beneficial but for more complex cases the difference would become apparent. \n",
+    "\n",
+    "A good option for a large model is gpt-4 but the inference time comparison will not be representative due to the hosting infrastructure differences. Try a larger model hosted by HF if you have access.\n",
+    "\n",
+    "The task here is to try out a large LLM and describe the differencies.\n",
+    "\n",
+    "The deliverables are:\n",
+    "\n",
+    "1. The comparison between outputs with Mistral-7b and a larger model\n",
+    "2. The difference in response times\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "942f39d4-eb27-4f2d-ae47-a5d65f102faa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Analysis of the difference between LLMs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "70c16440",
+   "metadata": {},
+   "source": [
+    "### Step 6 (Bonus): Use an LLM to quantitatively compare outputs of different variants of the system (LLM as a Judge)\n",
+    "\n",
+    "Use a powerful LLM (e.g. GPT-4) to quantitatively evaluate outputs of two alternative setups (different embedding models, different LLMs, both etc.). For inspiration and for prompts refer to [1](https://arxiv.org/pdf/2306.05685.pdf), [2](https://arxiv.org/pdf/2401.10020.pdf), [3](https://www.airtrain.ai/blog/the-comprehensive-guide-to-llm-evaluation#high-level-approach)\n",
+    "\n",
+    "The deliverables:\n",
+    "\n",
+    "1. The code you put together\n",
+    "2. The high-level description of the setup\n",
+    "3. The results of the qualitative comparison\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "39c18ba0-e54a-478f-9e60-0ea65c29238a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}