initial

README.md

@@ -0,0 +1,25 @@
+# A template for a RAG system with Gradio UI
+Deliberately stripped down to leave some room for experimenting
+
+# Setting it up
+- Clone https://github.com/huggingface/transformers to a local machine
+- Use the **prep_scrips/markdown_to_text.py** script to extract raw text from markdown from transformers/docs/source/en/
+- Break the resulting texts down into semantically meaningful pieces. Experiment with different chunking mechanisms to make sure the semantic meaning is captured.
+- Use **prep_scrips/lancedb_setup.py** to embed and store chunks in a lancedb instance. It also creates an index for fast ANN retrieval (not really needed for this exercise but necessary at scale). You'll need to put your own values into VECTOR_COLUMN_NAME, TEXT_COLUMN_NAME, DB_TABLE_NAME.
+- Move the database directory (.lancedb by default) to **gradio_app/**
+- Use the template given in **gradio_app** to wrap everything into the [Gradio](https://www.gradio.app/docs/interface) app and run it on HF [spaces](https://huggingface.co/docs/hub/spaces-config-reference). Make sure to adjust VECTOR_COLUMN_NAME, TEXT_COLUMN_NAME, DB_TABLE_NAME according to your DB setup.
+- Set up secrets OPENAI_API_KEY and HUGGING_FACE_HUB_TOKEN to use OpenAI and open-source models correspondingly
+
+- TODOs: 
+  - Experiment with chunking, see how it affects the results. When deciding how to chunk it helps to think about what kind of chunks you'd like to see as context to your queries.
+    - Deliverables: Show how retrieved documents differ with different chunking strategies and how it affects the output
+  - Try out different embedding models (EMB_MODEL_NAME). The models to try are **sentence-transformers/all-MiniLM-L6-v2** - lightweight, **thenlper/gte-large** - relatively heavy but more powerful.
+    - Deliverables: Show how retrieved documents differ with different embedding models and how they affect the output. Provide an estimate of how time to embed the chunks and DB ingestion time differs (happening in **prep_scrips/lancedb_setup.py**).
+  - Add a re-ranker (cross-encoder) to the pipeline. Start with sentence-transformers pages on cross-encoders [1](https://www.sbert.net/examples/applications/cross-encoder/README.html) [2](https://www.sbert.net/examples/applications/retrieve_rerank/README.html), then pick a [pretrained cross-encoder](https://www.sbert.net/docs/pretrained-models/ce-msmarco.html), e.g. **cross-encoder/ms-marco-MiniLM-L-12-v2**. Don't forget to increase the number of *retrieved* documents when using re-ranker. The number of document used as context should stay the same.
+    - Deliverables: Show how retrieved documents differ after adding a re-ranker and how it affects the output. Provide an estimate of how latency changes. 
+  - Try another LLM (e.g. LLaMA-2-70b, falcon-180b).
+    - Deliverables: Show how LLMs affect the output and how latency changes with the model size.
+  - Add more documents (e.g. diffusers, tokenizers, optimum etc) to see how the system scales.
+    - Deliverables: Show how latency changes, how it differs with and without index (index is added in **prep_scrips/lancedb_setup.py**).
+  - (Bonus) Use an LLM to quantitatively compare outputs of different variants of the system ([LLM as a Judge](https://huggingface.co/collections/andrewrreed/llm-as-a-judge-653fb861e361fd03c12d41e5))
+    - Deliverables: Describe the experimental setup and what were the evaluation results

gradio_app/app.py

@@ -0,0 +1,127 @@
+"""
+Credit to Derek Thomas, [email protected]
+"""
+
+import subprocess
+
+subprocess.run(["pip", "install", "--upgrade", "transformers[torch,sentencepiece]==4.34.1"])
+
+import logging
+from pathlib import Path
+from time import perf_counter
+
+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
+
+VECTOR_COLUMN_NAME = ""
+TEXT_COLUMN_NAME = ""
+
+proj_dir = Path(__file__).parent
+# Setting up the logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Set up the template environment with the templates directory
+env = Environment(loader=FileSystemLoader(proj_dir / 'templates'))
+
+# Load the templates directly from the environment
+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
+    history = history + [(text, None)]
+    return history, gr.Textbox(value="", interactive=False)
+
+
+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")
+
+    logger.warning('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]
+
+    document_time = perf_counter() - document_start
+    logger.warning(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
+
+    # Create Prompt
+    prompt = template.render(documents=documents, query=query)
+    prompt_html = template_html.render(documents=documents, query=query)
+
+    if api_kind == "HuggingFace":
+         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")
+
+    history[-1][1] = ""
+    for character in generate_fn(prompt, history[:-1]):
+        history[-1][1] = character
+        yield history, prompt_html
+
+
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot(
+            [],
+            elem_id="chatbot",
+            avatar_images=('https://aui.atlassian.com/aui/8.8/docs/images/avatar-person.svg',
+                           'https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg'),
+            bubble_full_width=False,
+            show_copy_button=True,
+            show_share_button=True,
+            )
+
+    with gr.Row():
+        txt = gr.Textbox(
+                scale=3,
+                show_label=False,
+                placeholder="Enter text and press enter",
+                container=False,
+                )
+        txt_btn = gr.Button(value="Submit text", scale=1)
+
+    api_kind = gr.Radio(choices=["HuggingFace", "OpenAI"], value="HuggingFace")
+
+    prompt_html = gr.HTML()
+    # Turn off interactivity while generating if you click
+    txt_msg = txt_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
+            bot, [chatbot, api_kind], [chatbot, prompt_html])
+
+    # Turn it back on
+    txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
+
+    # Turn off interactivity while generating if you hit enter
+    txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
+            bot, [chatbot, api_kind], [chatbot, prompt_html])
+
+    # 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

@@ -0,0 +1,154 @@
+import openai
+import gradio as gr
+
+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 = getenv("OPENAI_API_KEY")
+HF_TOKEN = getenv("HUGGING_FACE_HUB_TOKEN")
+
+hf_client = InferenceClient(
+        "mistralai/Mistral-7B-Instruct-v0.1",
+        token=HF_TOKEN
+        )
+
+
+def format_prompt(message: str, api_kind: str):
+    """
+    Formats the given message using a chat template.
+
+    Args:
+        message (str): The user message to be formatted.
+
+    Returns:
+        str: Formatted message after applying the chat template.
+    """
+
+    # Create a list of message dictionaries with role and content
+    messages: List[Dict[str, Any]] = [{'role': 'user', 'content': message}]
+
+    if api_kind == "openai":
+        return messages
+    elif api_kind == "hf":
+        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]:
+    """
+    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_new_tokens': max_new_tokens,
+        'top_p': top_p,
+        'repetition_penalty': repetition_penalty,
+        'do_sample': True,
+        'seed': 42,
+        }
+    
+    formatted_prompt = format_prompt(prompt, "hf")
+
+    try:
+        stream = hf_client.text_generation(formatted_prompt, **generate_kwargs,
+                                            stream=True, details=True, return_full_text=False)
+        output = ""
+        for response in stream:
+            output += response.token.text
+            yield output
+
+    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."
+        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"
+        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."
+
+
+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]:
+    """
+    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)
+        output = ""
+        for chunk in stream:
+            output += chunk.choices[0].delta.get("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."
+        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"
+        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."

gradio_app/backend/semantic_search.py

@@ -0,0 +1,18 @@
+import logging
+import lancedb
+import os
+from pathlib import Path
+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
+db_uri = os.path.join(Path(__file__).parents[1], ".lancedb")
+db = lancedb.connect(db_uri)
+table = db.open_table(DB_TABLE_NAME)

gradio_app/requirements.txt

@@ -0,0 +1,9 @@
+# transformers[torch,sentencepiece]==4.34.1
+wikiextractor==3.0.6
+sentence-transformers>2.2.0
+ipywidgets==8.1.1
+tqdm==4.66.1
+aiohttp==3.8.6
+huggingface-hub==0.17.3
+lancedb>=0.3
+openai==0.28

gradio_app/templates/template.j2

@@ -0,0 +1,8 @@
+Instructions: Use the following unique documents in the Context section to answer the Query at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
+Context:
+{% for doc in documents %}
+---
+    {{ doc }}
+{% endfor %}
+---
+Query: {{ query }}

gradio_app/templates/template_html.j2

@@ -0,0 +1,102 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Information Page</title>
+    <link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=Source+Sans+Pro:wght@400;600&amp;display=swap">
+    <link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;600&amp;display=swap">
+    <style>
+        * {
+            font-family: "Source Sans Pro";
+        }
+
+        .instructions > * {
+          color: #111 !important;
+        }
+
+        details.doc-box * {
+          color: #111 !important;
+        }
+
+        .dark {
+            background: #111;
+            color: white;
+        }
+
+        .doc-box {
+            padding: 10px;
+            margin-top: 10px;
+            background-color: #baecc2;
+            border-radius: 6px;
+            color: #111 !important;
+            max-width: 700px;
+            box-shadow: rgba(0, 0, 0, 0.2) 0px 1px 2px 0px;
+        }
+
+        .doc-full {
+            margin: 10px 14px;
+            line-height: 1.6rem;
+        }
+
+        .instructions {
+            color: #111 !important;
+            background: #b7bdfd;
+            display: block;
+            border-radius: 6px;
+            padding: 6px 10px;
+            line-height: 1.6rem;
+            max-width: 700px;
+            box-shadow: rgba(0, 0, 0, 0.2) 0px 1px 2px 0px;
+        }
+
+        .query {
+            color: #111 !important;
+            background: #ffbcbc;
+            display: block;
+            border-radius: 6px;
+            padding: 6px 10px;
+            line-height: 1.6rem;
+            max-width: 700px;
+            box-shadow: rgba(0, 0, 0, 0.2) 0px 1px 2px 0px;
+        }
+    </style>
+</head>
+<body>
+<div class="prose svelte-1ybaih5" id="component-6">
+<h2>Prompt</h2>
+Below is the prompt that is given to the model. <hr>
+<h2>Instructions</h2>
+    <span class="instructions">Use the following pieces of context to answer the question at the end.<br>If you don't know the answer, just say that you don't know, <span style="font-weight: bold;">don't try to make up an answer.</span></span><br>
+<h2>Context</h2>
+{% for doc in documents %}
+    <details class="doc-box">
+        <summary>
+            <b>Doc {{ loop.index }}:</b> <span class="doc-short">{{ doc[:100] }}...</span>
+        </summary>
+        <div class="doc-full">{{ doc }}</div>
+    </details>
+{% endfor %}
+
+    <h2>Query</h2>
+    <span class="query">{{ query }}</span>
+</div>
+
+<script>
+document.addEventListener("DOMContentLoaded", function() {
+    const detailsElements = document.querySelectorAll('.doc-box');
+
+    detailsElements.forEach(detail => {
+        detail.addEventListener('toggle', function() {
+            const docShort = this.querySelector('.doc-short');
+            if (this.open) {
+                docShort.style.display = 'none';
+            } else {
+                docShort.style.display = 'inline';
+            }
+        });
+    });
+});
+</script>
+</body>
+</html>

prep_scripts/lancedb_setup.py

@@ -0,0 +1,65 @@
+import lancedb
+import torch
+import pyarrow as pa
+import pandas as pd
+from pathlib import Path
+import tqdm
+import numpy as np
+
+from sentence_transformers import SentenceTransformer
+
+
+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)

prep_scripts/markdown_to_text.py

@@ -0,0 +1,49 @@
+from bs4 import BeautifulSoup
+from markdown import markdown
+import os
+import re
+from pathlib import Path
+
+
+DIR_TO_SCRAPE = "transformers/docs/source/en/"
+OUTPUT_DIR = str(Path().resolve() / "docs_dump")
+
+
+def markdown_to_text(markdown_string):
+    """ Converts a markdown string to plaintext """
+
+    # md -> html -> text since BeautifulSoup can extract text cleanly
+    html = markdown(markdown_string)
+
+    html = re.sub(r'<!--((.|\n)*)-->', '', html)
+    html = re.sub('<code>bash', '<code>', html)
+
+    # extract text
+    soup = BeautifulSoup(html, "html.parser")
+    text = ''.join(soup.findAll(text=True))
+
+    text = re.sub('```(py|diff|python)', '', text)
+    text = re.sub('```\n', '\n', text)
+    text = re.sub('-         .*', '', text)
+    text = text.replace('...', '')
+    text = re.sub('\n(\n)+', '\n\n', text)
+
+    return text
+
+
+dir_to_scrape = Path(DIR_TO_SCRAPE)
+files = list(dir_to_scrape.rglob("*"))
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+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)
+
+        with open(os.path.join(OUTPUT_DIR, f"{parent}_{file.stem}.txt"), "w") as f:
+            f.write(text)
+