Initial commit

.github/workflows/deploy_hf.yaml

@@ -0,0 +1,24 @@
+name: Sync to Hugging Face hub
+on:
+  push:
+    branches: [main]
+    paths:
+      - 'scripts/**'  # Replace with the path to your specific folder
+      - 'requirements.txt'  # Add this line to include requirements.txt
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          HF_USERNAME: ${{ secrets.HF_USERNAME }}
+        run: git push --force https://$HF_USERNAME:[email protected]/spaces/towardsai-buster/ai-tutor-chatbot main:main

.gitignore

@@ -0,0 +1,177 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+ai-tutor/
+venv_ai_tutor/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+.vscode/
+data/chroma-db**/
+.huggingface
+
+.DS_Store
+
+*.csv
+*.json
+*.jsonl
+*.html
+*.mdx
+*.pkl
+*.png
+*.mov

README.md

@@ -0,0 +1,74 @@
+---
+title: AI Tutor Chatbot
+emoji: 🧑🏻‍🏫
+colorFrom: gray
+colorTo: pink
+sdk: gradio
+sdk_version: 5.13.2
+app_file: scripts/main.py
+pinned: false
+---
+
+### Gradio UI Chatbot
+
+A Gradio UI for the chatbot is available in [scripts/main.py](./scripts/main.py).
+
+The Gradio demo is deployed on Hugging Face Spaces at: [AI Tutor Chatbot on Hugging Face](https://huggingface.co/spaces/towardsai-buster/ai-tutor-chatbot).
+
+**Note:** A GitHub Action automatically deploys the Gradio demo when changes are pushed to the `scripts` folder.
+
+### Installation (for Gradio UI)
+
+1. **Create a new Python environment:**
+
+   ```bash
+   python -m venv .venv
+   ```
+
+2. **Activate the environment:**
+
+   For macOS and Linux:
+
+   ```bash
+   source .venv/bin/activate
+   ```
+
+   For Windows:
+
+   ```bash
+   .venv\Scripts\activate
+   ```
+
+3. **Install the dependencies:**
+
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+### Usage (for Gradio UI)
+
+1. **Set environment variables:**
+
+   Before running the application, set up the required API keys:
+
+   For macOS and Linux:
+
+   ```bash
+   export OPENAI_API_KEY=your_openai_api_key_here
+   export COHERE_API_KEY=your_cohere_api_key_here
+   ```
+
+   For Windows:
+
+   ```bash
+   set OPENAI_API_KEY=your_openai_api_key_here
+   set COHERE_API_KEY=your_cohere_api_key_here
+   ```
+
+2. **Run the application:**
+
+   ```bash
+   python scripts/main.py
+   ```
+
+   This command starts the Gradio interface for the AI Tutor chatbot.

data/scraping_scripts/create_vector_stores.py

@@ -0,0 +1,218 @@
+"""
+Vector Store Creation Script
+
+Purpose:
+This script processes various data sources (e.g., transformers, peft, trl, llama_index, openai_cookbooks, langchain)
+to create vector stores using Chroma and LlamaIndex. It reads data from JSONL files, creates document embeddings,
+and stores them in persistent Chroma databases for efficient retrieval.
+
+Usage:
+python script_name.py <source1> <source2> ...
+
+Example:
+python script_name.py transformers peft llama_index
+
+The script accepts one or more source names as command-line arguments. Valid source names are:
+transformers, peft, trl, llama_index, openai_cookbooks, langchain
+
+For each specified source, the script will:
+1. Read data from the corresponding JSONL file
+2. Create document embeddings
+3. Store the embeddings in a Chroma vector database
+4. Save a dictionary of documents for future reference
+
+Note: Ensure that the input JSONL files are present in the 'data' directory.
+"""
+
+import argparse
+import json
+import os
+import pdb
+import pickle
+import shutil
+
+import chromadb
+from dotenv import load_dotenv
+from llama_index.core import Document, StorageContext, VectorStoreIndex
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.core.schema import MetadataMode, TextNode
+from llama_index.embeddings.cohere import CohereEmbedding
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.llms.openai import OpenAI
+from llama_index.vector_stores.chroma import ChromaVectorStore
+
+load_dotenv()
+
+# Configuration for different sources
+SOURCE_CONFIGS = {
+    "transformers": {
+        "input_file": "data/transformers_data.jsonl",
+        "db_name": "chroma-db-transformers",
+    },
+    "peft": {"input_file": "data/peft_data.jsonl", "db_name": "chroma-db-peft"},
+    "trl": {"input_file": "data/trl_data.jsonl", "db_name": "chroma-db-trl"},
+    "llama_index": {
+        "input_file": "data/llama_index_data.jsonl",
+        "db_name": "chroma-db-llama_index",
+    },
+    "openai_cookbooks": {
+        "input_file": "data/openai_cookbooks_data.jsonl",
+        "db_name": "chroma-db-openai_cookbooks",
+    },
+    "langchain": {
+        "input_file": "data/langchain_data.jsonl",
+        "db_name": "chroma-db-langchain",
+    },
+    "tai_blog": {
+        "input_file": "data/tai_blog_data.jsonl",
+        "db_name": "chroma-db-tai_blog",
+    },
+    "all_sources": {
+        "input_file": "data/all_sources_data.jsonl",
+        "db_name": "chroma-db-all_sources",
+    },
+}
+
+
+def create_docs(input_file: str) -> list[Document]:
+    with open(input_file, "r") as f:
+        documents = []
+        for line in f:
+            data = json.loads(line)
+            documents.append(
+                Document(
+                    doc_id=data["doc_id"],
+                    text=data["content"],
+                    metadata={  # type: ignore
+                        "url": data["url"],
+                        "title": data["name"],
+                        "tokens": data["tokens"],
+                        "retrieve_doc": data["retrieve_doc"],
+                        "source": data["source"],
+                    },
+                    excluded_llm_metadata_keys=[  # url is included in LLM context
+                        "title",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                    excluded_embed_metadata_keys=[  # title is embedded along the content
+                        "url",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                )
+            )
+    return documents
+
+
+def process_source(source: str):
+    config = SOURCE_CONFIGS[source]
+
+    input_file = config["input_file"]
+    db_name = config["db_name"]
+    db_path = f"data/{db_name}"
+
+    print(f"Processing source: {source}")
+
+    documents: list[Document] = create_docs(input_file)
+    print(f"Created {len(documents)} documents")
+
+    # Check if the folder exists and delete it
+    if os.path.exists(db_path):
+        print(f"Existing database found at {db_path}. Deleting...")
+        shutil.rmtree(db_path)
+        print(f"Deleted existing database at {db_path}")
+
+    # Create Chroma client and collection
+    chroma_client = chromadb.PersistentClient(path=f"data/{db_name}")
+    chroma_collection = chroma_client.create_collection(db_name)
+
+    # Create vector store and storage context
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
+
+    # Save document dictionary
+    document_dict: dict[str, Document] = {doc.doc_id: doc for doc in documents}
+    document_dict_file = f"data/{db_name}/document_dict_{source}.pkl"
+    with open(document_dict_file, "wb") as f:
+        pickle.dump(document_dict, f)
+    print(f"Saved document dictionary to {document_dict_file}")
+
+    # Load nodes with context
+    with open("data/all_sources_contextual_nodes.pkl", "rb") as f:
+        nodes_with_context: list[TextNode] = pickle.load(f)
+
+    print(f"Loaded {len(nodes_with_context)} nodes with context")
+    # pdb.set_trace()
+    # exit()
+
+    # Create vector store index
+    index = VectorStoreIndex(
+        nodes=nodes_with_context,
+        # embed_model=OpenAIEmbedding(model="text-embedding-3-large", mode="similarity"),
+        embed_model=CohereEmbedding(
+            api_key=os.environ["COHERE_API_KEY"],
+            model_name="embed-english-v3.0",
+            input_type="search_document",
+        ),
+        show_progress=True,
+        use_async=True,
+        storage_context=storage_context,
+    )
+    llm = OpenAI(
+        temperature=1,
+        model="gpt-4o-mini",
+        # model="gpt-4o",
+        max_tokens=5000,
+        max_retries=3,
+    )
+    query_engine = index.as_query_engine(llm=llm)
+    response = query_engine.query("How to fine-tune an llm?")
+    print(response)
+    for src in response.source_nodes:
+        print("Node ID\t", src.node_id)
+        print("Title\t", src.metadata["title"])
+        print("Text\t", src.text)
+        print("Score\t", src.score)
+        print("-_" * 20)
+
+    # # Create vector store index
+    # index = VectorStoreIndex.from_documents(
+    #     documents,
+    #     # embed_model=OpenAIEmbedding(model="text-embedding-3-large", mode="similarity"),
+    #     embed_model=CohereEmbedding(
+    #         api_key=os.environ["COHERE_API_KEY"],
+    #         model_name="embed-english-v3.0",
+    #         input_type="search_document",
+    #     ),
+    #     transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)],
+    #     show_progress=True,
+    #     use_async=True,
+    #     storage_context=storage_context,
+    # )
+    print(f"Created vector store index for {source}")
+
+
+def main(sources: list[str]):
+    for source in sources:
+        if source in SOURCE_CONFIGS:
+            process_source(source)
+        else:
+            print(f"Unknown source: {source}. Skipping.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Process sources and create vector stores."
+    )
+    parser.add_argument(
+        "sources",
+        nargs="+",
+        choices=SOURCE_CONFIGS.keys(),
+        help="Specify one or more sources to process",
+    )
+    args = parser.parse_args()
+
+    main(args.sources)

data/scraping_scripts/csv_to_jsonl.py

@@ -0,0 +1,61 @@
+import json
+import uuid
+
+import pandas as pd
+import tiktoken
+
+
+# Function to count tokens using tiktoken
+def num_tokens_from_string(string: str, encoding_name: str) -> int:
+    encoding = tiktoken.get_encoding(encoding_name)
+    num_tokens = len(
+        encoding.encode(
+            string, disallowed_special=(encoding.special_tokens_set - {"<|endoftext|>"})
+        )
+    )
+    return num_tokens
+
+
+# Function to clean or remove specific content, e.g., copyright headers
+def remove_copyright_header(content: str) -> str:
+    # Implement any cleaning logic you need here
+    return content
+
+
+# Function to convert DataFrame to JSONL format with token counting
+def convert_to_jsonl_with_conditions(df, encoding_name="cl100k_base"):
+    jsonl_data = []
+    for _, row in df.iterrows():
+        token_count = num_tokens_from_string(row["text"], encoding_name)
+
+        # Skip entries based on token count conditions
+        if token_count < 100 or token_count > 200_000:
+            print(f"Skipping {row['title']} due to token count {token_count}")
+            continue
+
+        cleaned_content = remove_copyright_header(row["text"])
+
+        entry = {
+            "tokens": token_count,  # Token count using tiktoken
+            "doc_id": str(uuid.uuid4()),  # Generate a unique UUID
+            "name": row["title"],
+            "url": row["tai_url"],
+            "retrieve_doc": (token_count <= 8000),  # retrieve_doc condition
+            "source": "tai_blog",
+            "content": cleaned_content,
+        }
+        jsonl_data.append(entry)
+    return jsonl_data
+
+
+# Load the CSV file
+data = pd.read_csv("data/tai.csv")
+
+# Convert the dataframe to JSONL format with token counting and conditions
+jsonl_data_with_conditions = convert_to_jsonl_with_conditions(data)
+
+# Save the output to a new JSONL file using json.dumps to ensure proper escaping
+output_path = "data/tai_blog_data_conditions.jsonl"
+with open(output_path, "w") as f:
+    for entry in jsonl_data_with_conditions:
+        f.write(json.dumps(entry) + "\n")

data/scraping_scripts/github_to_markdown_ai_docs.py

@@ -0,0 +1,214 @@
+"""
+Fetch Markdown files from specified GitHub repositories.
+
+This script fetches Markdown (.md), MDX (.mdx), and Jupyter Notebook (.ipynb) files
+from specified GitHub repositories, particularly focusing on documentation sources
+for various AI and machine learning libraries.
+
+Key features:
+1. Configurable for multiple documentation sources (e.g., Hugging Face Transformers, PEFT, TRL)
+2. Command-line interface for specifying one or more sources to process
+3. Automatic conversion of Jupyter Notebooks to Markdown
+4. Rate limiting handling to comply with GitHub API restrictions
+5. Retry mechanism for resilience against network issues
+
+Usage:
+    python github_to_markdown_ai_docs.py <source1> [<source2> ...]
+
+Where <sourceN> is one of the predefined sources in SOURCE_CONFIGS (e.g., 'transformers', 'peft', 'trl').
+
+Example:
+    python github_to_markdown_ai_docs.py trl peft
+
+This will download and process the documentation files for both TRL and PEFT libraries.
+
+Note: 
+- Ensure you have set the GITHUB_TOKEN variable with your GitHub Personal Access Token.
+- The script creates a 'data' directory in the current working directory to store the downloaded files.
+- Each source's files are stored in a subdirectory named '<repo>_md_files'.
+
+"""
+
+import argparse
+import json
+import os
+import random
+import time
+from typing import Dict, List
+
+import nbformat
+import requests
+from dotenv import load_dotenv
+from nbconvert import MarkdownExporter
+
+load_dotenv()
+
+# Configuration for different sources
+SOURCE_CONFIGS = {
+    "transformers": {
+        "owner": "huggingface",
+        "repo": "transformers",
+        "path": "docs/source/en",
+    },
+    "peft": {
+        "owner": "huggingface",
+        "repo": "peft",
+        "path": "docs/source",
+    },
+    "trl": {
+        "owner": "huggingface",
+        "repo": "trl",
+        "path": "docs/source",
+    },
+    "llama_index": {
+        "owner": "run-llama",
+        "repo": "llama_index",
+        "path": "docs/docs",
+    },
+    "openai_cookbooks": {
+        "owner": "openai",
+        "repo": "openai-cookbook",
+        "path": "examples",
+    },
+    "langchain": {
+        "owner": "langchain-ai",
+        "repo": "langchain",
+        "path": "docs/docs",
+    },
+}
+
+# GitHub Personal Access Token (replace with your own token)
+GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
+
+# Headers for authenticated requests
+HEADERS = {
+    "Authorization": f"token {GITHUB_TOKEN}",
+    "Accept": "application/vnd.github.v3+json",
+}
+
+# Maximum number of retries
+MAX_RETRIES = 5
+
+
+def check_rate_limit():
+    rate_limit_url = "https://api.github.com/rate_limit"
+    response = requests.get(rate_limit_url, headers=HEADERS)
+    data = response.json()
+    remaining = data["resources"]["core"]["remaining"]
+    reset_time = data["resources"]["core"]["reset"]
+
+    if remaining < 10:  # Adjust this threshold as needed
+        wait_time = reset_time - time.time()
+        print(f"Rate limit nearly exceeded. Waiting for {wait_time:.2f} seconds.")
+        time.sleep(wait_time + 1)  # Add 1 second buffer
+
+
+def get_files_in_directory(api_url: str, retries: int = 0) -> List[Dict]:
+    try:
+        check_rate_limit()
+        response = requests.get(api_url, headers=HEADERS)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        if retries < MAX_RETRIES:
+            wait_time = (2**retries) + random.random()
+            print(
+                f"Error fetching directory contents: {e}. Retrying in {wait_time:.2f} seconds..."
+            )
+            time.sleep(wait_time)
+            return get_files_in_directory(api_url, retries + 1)
+        else:
+            print(
+                f"Failed to fetch directory contents after {MAX_RETRIES} retries: {e}"
+            )
+            return []
+
+
+def download_file(file_url: str, file_path: str, retries: int = 0):
+    try:
+        check_rate_limit()
+        response = requests.get(file_url, headers=HEADERS)
+        response.raise_for_status()
+        with open(file_path, "wb") as file:
+            file.write(response.content)
+    except requests.exceptions.RequestException as e:
+        if retries < MAX_RETRIES:
+            wait_time = (2**retries) + random.random()
+            print(
+                f"Error downloading file: {e}. Retrying in {wait_time:.2f} seconds..."
+            )
+            time.sleep(wait_time)
+            download_file(file_url, file_path, retries + 1)
+        else:
+            print(f"Failed to download file after {MAX_RETRIES} retries: {e}")
+
+
+def convert_ipynb_to_md(ipynb_path: str, md_path: str):
+    with open(ipynb_path, "r", encoding="utf-8") as f:
+        notebook = nbformat.read(f, as_version=4)
+
+    exporter = MarkdownExporter()
+    markdown, _ = exporter.from_notebook_node(notebook)
+
+    with open(md_path, "w", encoding="utf-8") as f:
+        f.write(markdown)
+
+
+def fetch_files(api_url: str, local_dir: str):
+    files = get_files_in_directory(api_url)
+    for file in files:
+        if file["type"] == "file" and file["name"].endswith((".md", ".mdx", ".ipynb")):
+            file_url = file["download_url"]
+            file_name = file["name"]
+            file_path = os.path.join(local_dir, file_name)
+            print(f"Downloading {file_name}...")
+            download_file(file_url, file_path)
+
+            if file_name.endswith(".ipynb"):
+                md_file_name = file_name.replace(".ipynb", ".md")
+                md_file_path = os.path.join(local_dir, md_file_name)
+                print(f"Converting {file_name} to markdown...")
+                convert_ipynb_to_md(file_path, md_file_path)
+                os.remove(file_path)  # Remove the .ipynb file after conversion
+        elif file["type"] == "dir":
+            subdir = os.path.join(local_dir, file["name"])
+            os.makedirs(subdir, exist_ok=True)
+            fetch_files(file["url"], subdir)
+
+
+def process_source(source: str):
+    if source not in SOURCE_CONFIGS:
+        print(
+            f"Error: Unknown source '{source}'. Available sources: {', '.join(SOURCE_CONFIGS.keys())}"
+        )
+        return
+
+    config = SOURCE_CONFIGS[source]
+    api_url = f"https://api.github.com/repos/{config['owner']}/{config['repo']}/contents/{config['path']}"
+    local_dir = f"data/{config['repo']}_md_files"
+    os.makedirs(local_dir, exist_ok=True)
+
+    print(f"Processing source: {source}")
+    fetch_files(api_url, local_dir)
+    print(f"Finished processing {source}")
+
+
+def main(sources: List[str]):
+    for source in sources:
+        process_source(source)
+    print("All specified sources have been processed.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Fetch Markdown files from specified GitHub repositories."
+    )
+    parser.add_argument(
+        "sources",
+        nargs="+",
+        choices=SOURCE_CONFIGS.keys(),
+        help="Specify one or more sources to process",
+    )
+    args = parser.parse_args()
+
+    main(args.sources)

data/scraping_scripts/process_md_files.py

@@ -0,0 +1,579 @@
+"""
+Markdown Document Processor for Documentation Sources
+
+This script processes Markdown (.md) and MDX (.mdx) files from various documentation sources
+(such as Hugging Face Transformers, PEFT, TRL, LlamaIndex, and OpenAI Cookbook) and converts 
+them into a standardized JSONL format for further processing or indexing.
+
+Key features:
+1. Configurable for multiple documentation sources
+2. Extracts titles, generates URLs, and counts tokens for each document
+3. Supports inclusion/exclusion of specific directories and root files
+4. Removes copyright headers from content
+5. Generates a unique ID for each document
+6. Determines if a whole document should be retrieved based on token count
+7. Handles special cases like openai-cookbook repo by adding .ipynb extensions
+8. Processes multiple sources in a single run
+
+Usage:
+    python process_md_files.py <source1> <source2> ...
+
+Where <source1>, <source2>, etc. are one or more of the predefined sources in SOURCE_CONFIGS 
+(e.g., 'transformers', 'llama_index', 'openai_cookbooks').
+
+The script processes all Markdown files in the specified input directories (and their subdirectories),
+applies the configured filters, and saves the results in JSONL files. Each line in the output
+files represents a single document with metadata and content.
+
+To add or modify sources, update the SOURCE_CONFIGS dictionary at the top of the script.
+"""
+
+# import argparse
+# import json
+# import logging
+# import os
+# import re
+# import uuid
+# from typing import Dict, List
+
+# import tiktoken
+
+# logging.basicConfig(level=logging.INFO)
+# logger = logging.getLogger(__name__)
+
+# # Configuration for different sources
+# SOURCE_CONFIGS = {
+#     "transformers": {
+#         "base_url": "https://huggingface.co/docs/transformers/",
+#         "input_directory": "data/transformers_md_files",
+#         "output_file": "data/transformers_data.jsonl",
+#         "source_name": "transformers",
+#         "use_include_list": False,
+#         "included_dirs": [],
+#         "excluded_dirs": ["internal", "main_classes"],
+#         "excluded_root_files": [],
+#         "included_root_files": [],
+#         "url_extension": "",
+#     },
+#     "peft": {
+#         "base_url": "https://huggingface.co/docs/peft/",
+#         "input_directory": "data/peft_md_files",
+#         "output_file": "data/peft_data.jsonl",
+#         "source_name": "peft",
+#         "use_include_list": False,
+#         "included_dirs": [],
+#         "excluded_dirs": [],
+#         "excluded_root_files": [],
+#         "included_root_files": [],
+#         "url_extension": "",
+#     },
+#     "trl": {
+#         "base_url": "https://huggingface.co/docs/trl/",
+#         "input_directory": "data/trl_md_files",
+#         "output_file": "data/trl_data.jsonl",
+#         "source_name": "trl",
+#         "use_include_list": False,
+#         "included_dirs": [],
+#         "excluded_dirs": [],
+#         "excluded_root_files": [],
+#         "included_root_files": [],
+#         "url_extension": "",
+#     },
+#     "llama_index": {
+#         "base_url": "https://docs.llamaindex.ai/en/stable/",
+#         "input_directory": "data/llama_index_md_files",
+#         "output_file": "data/llama_index_data.jsonl",
+#         "source_name": "llama_index",
+#         "use_include_list": True,
+#         "included_dirs": [
+#             "getting_started",
+#             "understanding",
+#             "use_cases",
+#             "examples",
+#             "module_guides",
+#             "optimizing",
+#         ],
+#         "excluded_dirs": [],
+#         "excluded_root_files": [],
+#         "included_root_files": ["index.md"],
+#         "url_extension": "",
+#     },
+#     "openai_cookbooks": {
+#         "base_url": "https://github.com/openai/openai-cookbook/blob/main/examples/",
+#         "input_directory": "data/openai-cookbook_md_files",
+#         "output_file": "data/openai_cookbooks_data.jsonl",
+#         "source_name": "openai_cookbooks",
+#         "use_include_list": False,
+#         "included_dirs": [],
+#         "excluded_dirs": [],
+#         "excluded_root_files": [],
+#         "included_root_files": [],
+#         "url_extension": ".ipynb",
+#     },
+#     "langchain": {
+#         "base_url": "https://python.langchain.com/v0.2/docs/",
+#         "input_directory": "data/langchain_md_files",
+#         "output_file": "data/langchain_data.jsonl",
+#         "source_name": "langchain",
+#         "use_include_list": True,
+#         "included_dirs": ["how_to", "versions", "turorials", "integrations"],
+#         "excluded_dirs": [],
+#         "excluded_root_files": [],
+#         "included_root_files": ["security.md", "concepts.mdx", "introduction.mdx"],
+#         "url_extension": "",
+#     },
+#     "tai_blog": {
+#         "base_url": "",
+#         "input_directory": "",
+#         "output_file": "data/tai_blog_data.jsonl",
+#         "source_name": "tai_blog",
+#         "use_include_list": False,
+#         "included_dirs": [],
+#         "excluded_dirs": [],
+#         "excluded_root_files": [],
+#         "included_root_files": [],
+#         "url_extension": "",
+#     },
+# }
+
+
+# def extract_title(content: str):
+#     title_match = re.search(r"^#\s+(.+)$", content, re.MULTILINE)
+#     if title_match:
+#         return title_match.group(1).strip()
+
+#     lines = content.split("\n")
+#     for line in lines:
+#         if line.strip():
+#             return line.strip()
+
+#     return None
+
+
+# def generate_url(file_path: str, config: Dict) -> str:
+#     path_without_extension = os.path.splitext(file_path)[0]
+#     path_with_forward_slashes = path_without_extension.replace("\\", "/")
+#     return config["base_url"] + path_with_forward_slashes + config["url_extension"]
+
+
+# def should_include_file(file_path: str, config: Dict) -> bool:
+#     if os.path.dirname(file_path) == "":
+#         if config["use_include_list"]:
+#             return os.path.basename(file_path) in config["included_root_files"]
+#         else:
+#             return os.path.basename(file_path) not in config["excluded_root_files"]
+
+#     if config["use_include_list"]:
+#         return any(file_path.startswith(dir) for dir in config["included_dirs"])
+#     else:
+#         return not any(file_path.startswith(dir) for dir in config["excluded_dirs"])
+
+
+# def num_tokens_from_string(string: str, encoding_name: str) -> int:
+#     encoding = tiktoken.get_encoding(encoding_name)
+#     num_tokens = len(
+#         encoding.encode(
+#             string, disallowed_special=(encoding.special_tokens_set - {"<|endoftext|>"})
+#         )
+#     )
+#     return num_tokens
+
+
+# def remove_copyright_header(content: str) -> str:
+#     header_pattern = re.compile(r"<!--Copyright.*?-->\s*", re.DOTALL)
+#     cleaned_content = header_pattern.sub("", content, count=1)
+#     return cleaned_content.strip()
+
+
+# def process_md_files(directory: str, config: Dict) -> List[Dict]:
+#     jsonl_data = []
+
+#     for root, _, files in os.walk(directory):
+#         for file in files:
+#             if file.endswith(".md") or file.endswith(".mdx"):
+#                 file_path = os.path.join(root, file)
+#                 relative_path = os.path.relpath(file_path, directory)
+
+#                 if should_include_file(relative_path, config):
+#                     with open(file_path, "r", encoding="utf-8") as f:
+#                         content = f.read()
+
+#                     title = extract_title(content)
+#                     token_count = num_tokens_from_string(content, "cl100k_base")
+
+#                     if token_count < 100 or token_count > 200_000:
+#                         logger.info(
+#                             f"Skipping {relative_path} due to token count {token_count}"
+#                         )
+#                         continue
+
+#                     cleaned_content = remove_copyright_header(content)
+
+#                     json_object = {
+#                         "tokens": token_count,
+#                         "doc_id": str(uuid.uuid4()),
+#                         "name": (title if title else file),
+#                         "url": generate_url(relative_path, config),
+#                         "retrieve_doc": (token_count <= 8000),
+#                         "source": config["source_name"],
+#                         "content": cleaned_content,
+#                     }
+
+#                     jsonl_data.append(json_object)
+
+#     return jsonl_data
+
+
+# def save_jsonl(data: List[Dict], output_file: str) -> None:
+#     with open(output_file, "w", encoding="utf-8") as f:
+#         for item in data:
+#             json.dump(item, f, ensure_ascii=False)
+#             f.write("\n")
+
+
+# def combine_all_sources(sources: List[str]) -> None:
+#     all_data = []
+#     output_file = "data/all_sources_data.jsonl"
+
+#     for source in sources:
+#         if source not in SOURCE_CONFIGS:
+#             logger.error(f"Unknown source '{source}'. Skipping.")
+#             continue
+
+#         input_file = SOURCE_CONFIGS[source]["output_file"]
+#         logger.info(f"Processing source: {source}")
+
+#         with open(input_file, "r", encoding="utf-8") as f:
+#             for line in f:
+#                 all_data.append(json.loads(line))
+
+#     logger.info(f"Total documents combined: {len(all_data)}")
+#     save_jsonl(all_data, output_file)
+#     logger.info(f"Combined data saved to {output_file}")
+
+
+# def process_source(source: str) -> None:
+#     if source not in SOURCE_CONFIGS:
+#         logger.error(f"Unknown source '{source}'. Skipping.")
+#         return
+
+#     config = SOURCE_CONFIGS[source]
+#     logger.info(f"\n\nProcessing source: {source}")
+#     jsonl_data = process_md_files(config["input_directory"], config)
+#     save_jsonl(jsonl_data, config["output_file"])
+#     logger.info(
+#         f"Processed {len(jsonl_data)} files and saved to {config['output_file']}"
+#     )
+
+
+# def main(sources: List[str]) -> None:
+#     for source in sources:
+#         process_source(source)
+
+#     if len(sources) > 1:
+#         # sources = [
+#         #     "transformers",
+#         #     "peft",
+#         #     "trl",
+#         #     "llama_index",
+#         #     "langchain",
+#         #     "openai_cookbooks",
+#         #     "tai_blog",
+#         # ]
+#         combine_all_sources(sources)
+
+
+# if __name__ == "__main__":
+#     parser = argparse.ArgumentParser(
+#         description="Process Markdown files from specified sources."
+#     )
+#     parser.add_argument(
+#         "sources",
+#         nargs="+",
+#         choices=SOURCE_CONFIGS.keys(),
+#         help="Specify one or more sources to process",
+#     )
+#     args = parser.parse_args()
+
+#     main(args.sources)
+
+
+import argparse
+import json
+import logging
+import os
+import re
+import uuid
+from typing import Dict, List
+
+import tiktoken
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Configuration for different sources
+SOURCE_CONFIGS = {
+    "transformers": {
+        "base_url": "https://huggingface.co/docs/transformers/",
+        "input_directory": "data/transformers_md_files",
+        "output_file": "data/transformers_data.jsonl",
+        "source_name": "transformers",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": ["internal", "main_classes"],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "peft": {
+        "base_url": "https://huggingface.co/docs/peft/",
+        "input_directory": "data/peft_md_files",
+        "output_file": "data/peft_data.jsonl",
+        "source_name": "peft",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "trl": {
+        "base_url": "https://huggingface.co/docs/trl/",
+        "input_directory": "data/trl_md_files",
+        "output_file": "data/trl_data.jsonl",
+        "source_name": "trl",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "llama_index": {
+        "base_url": "https://docs.llamaindex.ai/en/stable/",
+        "input_directory": "data/llama_index_md_files",
+        "output_file": "data/llama_index_data.jsonl",
+        "source_name": "llama_index",
+        "use_include_list": True,
+        "included_dirs": [
+            "getting_started",
+            "understanding",
+            "use_cases",
+            "examples",
+            "module_guides",
+            "optimizing",
+        ],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": ["index.md"],
+        "url_extension": "",
+    },
+    "openai_cookbooks": {
+        "base_url": "https://github.com/openai/openai-cookbook/blob/main/examples/",
+        "input_directory": "data/openai-cookbook_md_files",
+        "output_file": "data/openai_cookbooks_data.jsonl",
+        "source_name": "openai_cookbooks",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": ".ipynb",
+    },
+    "langchain": {
+        "base_url": "https://python.langchain.com/v0.2/docs/",
+        "input_directory": "data/langchain_md_files",
+        "output_file": "data/langchain_data.jsonl",
+        "source_name": "langchain",
+        "use_include_list": True,
+        "included_dirs": ["how_to", "versions", "turorials", "integrations"],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": ["security.md", "concepts.mdx", "introduction.mdx"],
+        "url_extension": "",
+    },
+    "tai_blog": {
+        "base_url": "",
+        "input_directory": "",
+        "output_file": "data/tai_blog_data.jsonl",
+        "source_name": "tai_blog",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+    "8-hour_primer": {
+        "base_url": "",
+        "input_directory": "data/8-hour_primer",
+        "output_file": "data/8-hour_primer_data.jsonl",  # 8-hour Generative AI Primer
+        "source_name": "8-hour_primer",
+        "use_include_list": False,
+        "included_dirs": [],
+        "excluded_dirs": [],
+        "excluded_root_files": [],
+        "included_root_files": [],
+        "url_extension": "",
+    },
+}
+
+
+def extract_title(content: str):
+    title_match = re.search(r"^#\s+(.+)$", content, re.MULTILINE)
+    if title_match:
+        return title_match.group(1).strip()
+
+    lines = content.split("\n")
+    for line in lines:
+        if line.strip():
+            return line.strip()
+
+    return None
+
+
+def generate_url(file_path: str, config: Dict) -> str:
+    """
+    Return an empty string if base_url is empty;
+    otherwise return the constructed URL as before.
+    """
+    if not config["base_url"]:
+        return ""
+
+    path_without_extension = os.path.splitext(file_path)[0]
+    path_with_forward_slashes = path_without_extension.replace("\\", "/")
+    return config["base_url"] + path_with_forward_slashes + config["url_extension"]
+
+
+def should_include_file(file_path: str, config: Dict) -> bool:
+    if os.path.dirname(file_path) == "":
+        if config["use_include_list"]:
+            return os.path.basename(file_path) in config["included_root_files"]
+        else:
+            return os.path.basename(file_path) not in config["excluded_root_files"]
+
+    if config["use_include_list"]:
+        return any(file_path.startswith(dir) for dir in config["included_dirs"])
+    else:
+        return not any(file_path.startswith(dir) for dir in config["excluded_dirs"])
+
+
+def num_tokens_from_string(string: str, encoding_name: str) -> int:
+    encoding = tiktoken.get_encoding(encoding_name)
+    num_tokens = len(
+        encoding.encode(
+            string, disallowed_special=(encoding.special_tokens_set - {"<|endoftext|>"})
+        )
+    )
+    return num_tokens
+
+
+def remove_copyright_header(content: str) -> str:
+    header_pattern = re.compile(r"<!--Copyright.*?-->\s*", re.DOTALL)
+    cleaned_content = header_pattern.sub("", content, count=1)
+    return cleaned_content.strip()
+
+
+def process_md_files(directory: str, config: Dict) -> List[Dict]:
+    jsonl_data = []
+
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if file.endswith(".md") or file.endswith(".mdx"):
+                file_path = os.path.join(root, file)
+                relative_path = os.path.relpath(file_path, directory)
+
+                if should_include_file(relative_path, config):
+                    with open(file_path, "r", encoding="utf-8") as f:
+                        content = f.read()
+
+                    title = extract_title(content)
+                    token_count = num_tokens_from_string(content, "cl100k_base")
+
+                    # Skip very small or extremely large files
+                    if token_count < 100 or token_count > 200_000:
+                        logger.info(
+                            f"Skipping {relative_path} due to token count {token_count}"
+                        )
+                        continue
+
+                    cleaned_content = remove_copyright_header(content)
+
+                    json_object = {
+                        "tokens": token_count,
+                        "doc_id": str(uuid.uuid4()),
+                        "name": (title if title else file),
+                        "url": generate_url(relative_path, config),
+                        "retrieve_doc": (token_count <= 8000),
+                        "source": config["source_name"],
+                        "content": cleaned_content,
+                    }
+
+                    jsonl_data.append(json_object)
+
+    return jsonl_data
+
+
+def save_jsonl(data: List[Dict], output_file: str) -> None:
+    with open(output_file, "w", encoding="utf-8") as f:
+        for item in data:
+            json.dump(item, f, ensure_ascii=False)
+            f.write("\n")
+
+
+def combine_all_sources(sources: List[str]) -> None:
+    all_data = []
+    output_file = "data/all_sources_data.jsonl"
+
+    for source in sources:
+        if source not in SOURCE_CONFIGS:
+            logger.error(f"Unknown source '{source}'. Skipping.")
+            continue
+
+        input_file = SOURCE_CONFIGS[source]["output_file"]
+        logger.info(f"Processing source: {source}")
+
+        with open(input_file, "r", encoding="utf-8") as f:
+            for line in f:
+                all_data.append(json.loads(line))
+
+    logger.info(f"Total documents combined: {len(all_data)}")
+    save_jsonl(all_data, output_file)
+    logger.info(f"Combined data saved to {output_file}")
+
+
+def process_source(source: str) -> None:
+    if source not in SOURCE_CONFIGS:
+        logger.error(f"Unknown source '{source}'. Skipping.")
+        return
+
+    config = SOURCE_CONFIGS[source]
+    logger.info(f"\n\nProcessing source: {source}")
+    jsonl_data = process_md_files(config["input_directory"], config)
+    save_jsonl(jsonl_data, config["output_file"])
+    logger.info(
+        f"Processed {len(jsonl_data)} files and saved to {config['output_file']}"
+    )
+
+
+def main(sources: List[str]) -> None:
+    for source in sources:
+        process_source(source)
+
+    if len(sources) > 1:
+        combine_all_sources(sources)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Process Markdown files from specified sources."
+    )
+    parser.add_argument(
+        "sources",
+        nargs="+",
+        choices=SOURCE_CONFIGS.keys(),
+        help="Specify one or more sources to process",
+    )
+    args = parser.parse_args()
+
+    main(args.sources)

data/scraping_scripts/upload_dbs_to_hf.py

@@ -0,0 +1,38 @@
+"""
+Hugging Face Data Upload Script
+
+Purpose:
+This script uploads a local folder to a Hugging Face dataset repository. It's designed to
+update or create a dataset on the Hugging Face Hub by uploading the contents of a specified
+local folder.
+
+Usage:
+- Run the script: python data/scraping_scripts/upload_dbs_to_hf.py
+
+The script will:
+- Upload the contents of the 'data' folder to the specified Hugging Face dataset repository.
+- https://huggingface.co/datasets/towardsai-buster/ai-tutor-vector-db
+
+Configuration:
+- The script is set to upload to the "towardsai-buster/test-data" dataset repository. 
+- It deletes all existing files in the repository before uploading (due to delete_patterns=["*"]).
+"""
+
+import os
+
+from dotenv import load_dotenv
+from huggingface_hub import HfApi
+
+load_dotenv()
+
+api = HfApi(token=os.getenv("HF_TOKEN"))
+
+api.upload_folder(
+    folder_path="data",
+    repo_id="towardsai-tutors/ai-tutor-vector-db",
+    repo_type="dataset",
+    # multi_commits=True,
+    # multi_commits_verbose=True,
+    delete_patterns=["*"],
+    ignore_patterns=["*.jsonl", "*.py", "*.txt", "*.ipynb", "*.md", "*.pyc"],
+)

requirements.txt

@@ -0,0 +1,20 @@
+modal
+openai
+anthropic
+instructor
+pydantic
+logfire
+chromadb
+cohere
+tiktoken
+llama-index
+llama-index-postprocessor-cohere-rerank
+llama-index-embeddings-cohere
+llama-index-vector-stores-chroma
+python-dotenv
+ipykernel
+google-generativeai
+llama-index-llms-gemini
+gradio
+pymongo
+huggingface_hub

scripts/contextual_retrieval.py

@@ -0,0 +1,244 @@
+import asyncio
+import json
+import pdb
+import pickle
+from typing import Dict, List
+
+import instructor
+import logfire
+import tiktoken
+from anthropic import AsyncAnthropic
+from dotenv import load_dotenv
+
+# from instructor import AsyncInstructor, Mode, patch
+from jinja2 import Template
+from llama_index.core import Document
+from llama_index.core.ingestion import IngestionPipeline
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.core.schema import TextNode
+from openai import AsyncOpenAI
+from pydantic import BaseModel, Field
+from tenacity import retry, stop_after_attempt, wait_exponential
+from tqdm.asyncio import tqdm
+
+load_dotenv(".env")
+
+# logfire.configure()
+
+
+def create_docs(input_file: str) -> List[Document]:
+    with open(input_file, "r") as f:
+        documents: list[Document] = []
+        for line in f:
+            data = json.loads(line)
+            documents.append(
+                Document(
+                    doc_id=data["doc_id"],
+                    text=data["content"],
+                    metadata={  # type: ignore
+                        "url": data["url"],
+                        "title": data["name"],
+                        "tokens": data["tokens"],
+                        "retrieve_doc": data["retrieve_doc"],
+                        "source": data["source"],
+                    },
+                    excluded_llm_metadata_keys=[
+                        "title",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                    excluded_embed_metadata_keys=[
+                        "url",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                )
+            )
+    return documents
+
+
+class SituatedContext(BaseModel):
+    title: str = Field(..., description="The title of the document.")
+    context: str = Field(
+        ..., description="The context to situate the chunk within the document."
+    )
+
+
+# client = AsyncInstructor(
+#     client=AsyncAnthropic(),
+#     create=patch(
+#         create=AsyncAnthropic().beta.prompt_caching.messages.create,
+#         mode=Mode.ANTHROPIC_TOOLS,
+#     ),
+#     mode=Mode.ANTHROPIC_TOOLS,
+# )
+aclient = AsyncOpenAI()
+# logfire.instrument_openai(aclient)
+client: instructor.AsyncInstructor = instructor.from_openai(aclient)
+
+
+@retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=10))
+async def situate_context(doc: str, chunk: str) -> str:
+    template = Template(
+        """
+<document>
+{{ doc }}
+</document>
+
+Here is the chunk we want to situate within the whole document above:
+
+<chunk>
+{{ chunk }}
+</chunk>
+
+Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk.
+Answer only with the succinct context and nothing else.
+"""
+    )
+
+    content = template.render(doc=doc, chunk=chunk)
+
+    response = await client.chat.completions.create(
+        model="gpt-4o-mini",
+        max_tokens=1000,
+        temperature=0,
+        messages=[
+            {
+                "role": "user",
+                "content": content,
+            }
+        ],
+        response_model=SituatedContext,
+    )
+    return response.context
+
+
+# async def process_chunk(node: TextNode, document_dict: dict) -> TextNode:
+
+#     doc_id: str = node.source_node.node_id  # type: ignore
+#     doc: Document = document_dict[doc_id]
+
+#     if doc.metadata["tokens"] > 120_000:
+#         # Tokenize the document text
+#         encoding = tiktoken.encoding_for_model("gpt-4o-mini")
+#         tokens = encoding.encode(doc.text)
+
+#         # Trim to 120,000 tokens
+#         trimmed_tokens = tokens[:120_000]
+
+#         # Decode back to text
+#         trimmed_text = encoding.decode(trimmed_tokens)
+
+#         # Update the document text
+#         doc.text = trimmed_text
+#         doc.metadata["tokens"] = 120_000
+
+#     context: str = await situate_context(doc.text, node.text)
+#     node.text = f"{node.text}\n\n{context}"
+#     return node
+
+
+async def process_chunk(node: TextNode, document_dict: dict) -> TextNode:
+    doc_id: str = node.source_node.node_id  # type: ignore
+    doc: Document = document_dict[doc_id]
+
+    if doc.metadata["tokens"] > 120_000:
+        # Tokenize the document text
+        encoding = tiktoken.encoding_for_model("gpt-4o-mini")
+        tokens = encoding.encode(doc.get_content())
+
+        # Trim to 120,000 tokens
+        trimmed_tokens = tokens[:120_000]
+
+        # Decode back to text
+        trimmed_text = encoding.decode(trimmed_tokens)
+
+        # Update the document with trimmed text
+        doc = Document(text=trimmed_text, metadata=doc.metadata)
+        doc.metadata["tokens"] = 120_000
+
+    context: str = await situate_context(doc.get_content(), node.text)
+    node.text = f"{node.text}\n\n{context}"
+    return node
+
+
+async def process(
+    documents: List[Document], semaphore_limit: int = 50
+) -> List[TextNode]:
+
+    # From the document, we create chunks
+    pipeline = IngestionPipeline(
+        transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)]
+    )
+    all_nodes: list[TextNode] = pipeline.run(documents=documents, show_progress=True)
+    print(f"Number of nodes: {len(all_nodes)}")
+
+    document_dict: dict[str, Document] = {doc.doc_id: doc for doc in documents}
+
+    semaphore = asyncio.Semaphore(semaphore_limit)
+
+    async def process_with_semaphore(node):
+        async with semaphore:
+            result = await process_chunk(node, document_dict)
+            await asyncio.sleep(0.1)
+            return result
+
+    tasks = [process_with_semaphore(node) for node in all_nodes]
+
+    results: List[TextNode] = await tqdm.gather(*tasks, desc="Processing chunks")
+
+    # results: List[TextNode] = []
+    # # Add tqdm progress bar with semaphore limit
+    # for task in tqdm(
+    #     asyncio.as_completed(tasks), total=len(tasks), desc="Processing chunks"
+    # ):
+    #     result = await task
+    #     results.append(result)
+    # pdb.set_trace()
+
+    return results
+
+
+async def main():
+    documents: List[Document] = create_docs("data/all_sources_data.jsonl")
+    enhanced_nodes: List[TextNode] = await process(documents)
+
+    with open("data/all_sources_contextual_nodes.pkl", "wb") as f:
+        pickle.dump(enhanced_nodes, f)
+
+    # pipeline = IngestionPipeline(
+    #     transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)]
+    # )
+    # all_nodes: list[TextNode] = pipeline.run(documents=documents, show_progress=True)
+    # print(all_nodes[7933])
+    # pdb.set_trace()
+
+    with open("data/all_sources_contextual_nodes.pkl", "rb") as f:
+        enhanced_nodes: list[TextNode] = pickle.load(f)
+
+    for i, node in enumerate(enhanced_nodes):
+        print(f"Chunk {i + 1}:")
+        print(f"Node: {node}")
+        print(f"Text: {node.text}")
+        # pdb.set_trace()
+        break
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+
+
+# Ok so I need to create a new chroma-db-all_sources that embedded (context+chunk)
+# I need to create an index and instead of from_documents it will be from_nodes
+
+
+# First I need to create contexts for each chunk. Create a list of tasks (doc + chunk)
+
+
+# documents = create_docs("data/all_sources_data.jsonl")
+# pipeline = IngestionPipeline(
+#     transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)]
+# )
+# all_nodes = pipeline.run(documents=documents, show_progress=True)

scripts/custom_retriever.py

@@ -0,0 +1,257 @@
+import asyncio
+import os
+import time
+import traceback
+from typing import List, Optional
+
+import logfire
+import tiktoken
+from cohere import AsyncClient
+from dotenv import load_dotenv
+from llama_index.core import Document, QueryBundle
+from llama_index.core.async_utils import run_async_tasks
+from llama_index.core.callbacks import CBEventType, EventPayload
+from llama_index.core.retrievers import (
+    BaseRetriever,
+    KeywordTableSimpleRetriever,
+    VectorIndexRetriever,
+)
+from llama_index.core.schema import MetadataMode, NodeWithScore, QueryBundle, TextNode
+from llama_index.core.vector_stores import (
+    FilterCondition,
+    FilterOperator,
+    MetadataFilter,
+    MetadataFilters,
+)
+from llama_index.postprocessor.cohere_rerank import CohereRerank
+from llama_index.postprocessor.cohere_rerank.base import CohereRerank
+
+load_dotenv()
+
+
+class AsyncCohereRerank(CohereRerank):
+    def __init__(
+        self,
+        top_n: int = 5,
+        model: str = "rerank-english-v3.0",
+        api_key: Optional[str] = None,
+    ) -> None:
+        super().__init__(top_n=top_n, model=model, api_key=api_key)
+        self._api_key = api_key
+        self._model = model
+        self._top_n = top_n
+
+    async def apostprocess_nodes(
+        self,
+        nodes: List[NodeWithScore],
+        query_bundle: Optional[QueryBundle] = None,
+    ) -> List[NodeWithScore]:
+        if query_bundle is None:
+            raise ValueError("Query bundle must be provided.")
+
+        if len(nodes) == 0:
+            return []
+
+        async_client = AsyncClient(api_key=self._api_key)
+
+        with self.callback_manager.event(
+            CBEventType.RERANKING,
+            payload={
+                EventPayload.NODES: nodes,
+                EventPayload.MODEL_NAME: self._model,
+                EventPayload.QUERY_STR: query_bundle.query_str,
+                EventPayload.TOP_K: self._top_n,
+            },
+        ) as event:
+            texts = [
+                node.node.get_content(metadata_mode=MetadataMode.EMBED)
+                for node in nodes
+            ]
+
+            results = await async_client.rerank(
+                model=self._model,
+                top_n=self._top_n,
+                query=query_bundle.query_str,
+                documents=texts,
+            )
+
+            new_nodes = []
+            for result in results.results:
+                new_node_with_score = NodeWithScore(
+                    node=nodes[result.index].node, score=result.relevance_score
+                )
+                new_nodes.append(new_node_with_score)
+            event.on_end(payload={EventPayload.NODES: new_nodes})
+
+        return new_nodes
+
+
+class CustomRetriever(BaseRetriever):
+    """Custom retriever that performs both semantic search and hybrid search."""
+
+    def __init__(
+        self,
+        vector_retriever: VectorIndexRetriever,
+        document_dict: dict,
+        keyword_retriever=None,
+        mode: str = "AND",
+    ) -> None:
+        """Init params."""
+        self._vector_retriever = vector_retriever
+        self._document_dict = document_dict
+        self._keyword_retriever = keyword_retriever
+        if mode not in ("AND", "OR"):
+            raise ValueError("Invalid mode.")
+        self._mode = mode
+        super().__init__()
+
+    async def _process_retrieval(
+        self, query_bundle: QueryBundle, is_async: bool = True
+    ) -> List[NodeWithScore]:
+        """Common processing logic for both sync and async retrieval."""
+        # Clean query string
+        query_bundle.query_str = query_bundle.query_str.replace(
+            "\ninput is ", ""
+        ).rstrip()
+        logfire.info(f"Retrieving nodes with string: '{query_bundle}'")
+
+        start = time.time()
+
+        # Get nodes from both retrievers
+        if is_async:
+            nodes = await self._vector_retriever.aretrieve(query_bundle)
+        else:
+            nodes = self._vector_retriever.retrieve(query_bundle)
+
+        keyword_nodes = []
+        if self._keyword_retriever:
+            if is_async:
+                keyword_nodes = await self._keyword_retriever.aretrieve(query_bundle)
+            else:
+                keyword_nodes = self._keyword_retriever.retrieve(query_bundle)
+
+        logfire.info(f"Number of vector nodes: {len(nodes)}")
+        logfire.info(f"Number of keyword nodes: {len(keyword_nodes)}")
+
+        # # Filter keyword nodes based on metadata filters from vector retriever
+        # if (
+        #     hasattr(self._vector_retriever, "_filters")
+        #     and self._vector_retriever._filters
+        # ):
+        #     filtered_keyword_nodes = []
+        #     for node in keyword_nodes:
+        #         node_source = node.node.metadata.get("source")
+        #         # Check if node's source matches any of the filter conditions
+        #         for filter in self._vector_retriever._filters.filters:
+        #             if (
+        #                 isinstance(filter, MetadataFilter)
+        #                 and filter.key == "source"
+        #                 and filter.operator == FilterOperator.EQ
+        #                 and filter.value == node_source
+        #             ):
+        #                 filtered_keyword_nodes.append(node)
+        #                 break
+        #     keyword_nodes = filtered_keyword_nodes
+        #     logfire.info(
+        #         f"Number of keyword nodes after filtering: {len(keyword_nodes)}"
+        #     )
+
+        # Combine results based on mode
+        vector_ids = {n.node.node_id for n in nodes}
+        keyword_ids = {n.node.node_id for n in keyword_nodes}
+        combined_dict = {n.node.node_id: n for n in nodes}
+        combined_dict.update({n.node.node_id: n for n in keyword_nodes})
+
+        # If no keyword retriever or no keyword nodes, just use vector nodes
+        if not self._keyword_retriever or not keyword_nodes:
+            retrieve_ids = vector_ids
+        else:
+            retrieve_ids = (
+                vector_ids.intersection(keyword_ids)
+                if self._mode == "AND"
+                else vector_ids.union(keyword_ids)
+            )
+
+        nodes = [combined_dict[rid] for rid in retrieve_ids]
+        logfire.info(f"Number of combined nodes: {len(nodes)}")
+
+        # Filter unique doc IDs
+        nodes = self._filter_nodes_by_unique_doc_id(nodes)
+        logfire.info(f"Number of nodes without duplicate doc IDs: {len(nodes)}")
+
+        # Process node contents
+        for node in nodes:
+            doc_id = node.node.source_node.node_id
+            if node.metadata["retrieve_doc"]:
+                doc = self._document_dict[doc_id]
+                node.node.text = doc.text
+            node.node.node_id = doc_id
+
+        # Rerank results
+        try:
+            reranker = (
+                AsyncCohereRerank(top_n=5, model="rerank-english-v3.0")
+                if is_async
+                else CohereRerank(top_n=5, model="rerank-english-v3.0")
+            )
+            nodes = (
+                await reranker.apostprocess_nodes(nodes, query_bundle)
+                if is_async
+                else reranker.postprocess_nodes(nodes, query_bundle)
+            )
+        except Exception as e:
+            error_msg = f"Error during reranking: {type(e).__name__}: {str(e)}\n"
+            error_msg += "Traceback:\n"
+            error_msg += traceback.format_exc()
+            logfire.error(error_msg)
+
+        # Filter by score and token count
+        nodes_filtered = self._filter_by_score_and_tokens(nodes)
+
+        duration = time.time() - start
+        logfire.info(f"Retrieving nodes took {duration:.2f}s")
+        logfire.info(f"Nodes sent to LLM: {nodes_filtered[:5]}")
+
+        return nodes_filtered[:5]
+
+    def _filter_nodes_by_unique_doc_id(
+        self, nodes: List[NodeWithScore]
+    ) -> List[NodeWithScore]:
+        """Filter nodes to keep only unique doc IDs."""
+        unique_nodes = {}
+        for node in nodes:
+            doc_id = node.node.source_node.node_id
+            if doc_id is not None and doc_id not in unique_nodes:
+                unique_nodes[doc_id] = node
+        return list(unique_nodes.values())
+
+    def _filter_by_score_and_tokens(
+        self, nodes: List[NodeWithScore]
+    ) -> List[NodeWithScore]:
+        """Filter nodes by score and token count."""
+        nodes_filtered = []
+        total_tokens = 0
+        enc = tiktoken.encoding_for_model("gpt-4o-mini")
+
+        for node in nodes:
+            if node.score < 0.10:
+                logfire.info(f"Skipping node with score {node.score}")
+                continue
+
+            node_tokens = len(enc.encode(node.node.text))
+            if total_tokens + node_tokens > 100_000:
+                logfire.info("Skipping node due to token count exceeding 100k")
+                break
+
+            total_tokens += node_tokens
+            nodes_filtered.append(node)
+
+        return nodes_filtered
+
+    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Async retrieve nodes given query."""
+        return await self._process_retrieval(query_bundle, is_async=True)
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Sync retrieve nodes given query."""
+        return asyncio.run(self._process_retrieval(query_bundle, is_async=False))

scripts/evaluate_rag_system.py

@@ -0,0 +1,773 @@
+import asyncio
+import html
+import json
+import logging
+import os
+import pdb
+import pickle
+import random
+import time
+from typing import Dict, List, Optional, Tuple
+
+import aiofiles
+import chromadb
+import logfire
+import pandas as pd
+from custom_retriever import CustomRetriever
+from llama_index.agent.openai import OpenAIAgent
+from llama_index.core import Document, SimpleKeywordTableIndex, VectorStoreIndex
+from llama_index.core.base.base_retriever import BaseRetriever
+from llama_index.core.bridge.pydantic import Field, SerializeAsAny
+from llama_index.core.chat_engine.types import (
+    AGENT_CHAT_RESPONSE_TYPE,
+    AgentChatResponse,
+    ChatResponseMode,
+)
+from llama_index.core.evaluation import (
+    AnswerRelevancyEvaluator,
+    BatchEvalRunner,
+    EmbeddingQAFinetuneDataset,
+    FaithfulnessEvaluator,
+    RelevancyEvaluator,
+)
+from llama_index.core.evaluation.base import EvaluationResult
+from llama_index.core.evaluation.retrieval.base import (
+    BaseRetrievalEvaluator,
+    RetrievalEvalMode,
+    RetrievalEvalResult,
+)
+from llama_index.core.indices.base_retriever import BaseRetriever
+from llama_index.core.ingestion import IngestionPipeline
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+from llama_index.core.retrievers import (
+    BaseRetriever,
+    KeywordTableSimpleRetriever,
+    VectorIndexRetriever,
+)
+from llama_index.core.schema import ImageNode, NodeWithScore, QueryBundle, TextNode
+from llama_index.core.tools import RetrieverTool, ToolMetadata
+from llama_index.core.vector_stores import (
+    FilterOperator,
+    MetadataFilter,
+    MetadataFilters,
+)
+from llama_index.embeddings.cohere import CohereEmbedding
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.llms.gemini import Gemini
+from llama_index.llms.openai import OpenAI
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from prompts import system_message_openai_agent
+from pydantic import BaseModel, Field
+from tqdm.asyncio import tqdm_asyncio
+
+# from setup import (
+#     AVAILABLE_SOURCES,
+#     AVAILABLE_SOURCES_UI,
+#     custom_retriever_all_sources,
+#     custom_retriever_langchain,
+#     custom_retriever_llama_index,
+#     custom_retriever_openai_cookbooks,
+#     custom_retriever_peft,
+#     custom_retriever_transformers,
+#     custom_retriever_trl,
+# )
+
+
+class RotatingJSONLWriter:
+    def __init__(
+        self, base_filename: str, max_size: int = 10**6, backup_count: int = 5
+    ):
+        """
+        Initialize the rotating JSONL writer.
+
+        Args:
+            base_filename (str): The base filename for the JSONL files.
+            max_size (int): Maximum size in bytes before rotating.
+            backup_count (int): Number of backup files to keep.
+        """
+        self.base_filename = base_filename
+        self.max_size = max_size
+        self.backup_count = backup_count
+        self.current_file = base_filename
+
+    async def write(self, data: dict):
+        # Rotate if file exceeds max size
+        if (
+            os.path.exists(self.current_file)
+            and os.path.getsize(self.current_file) > self.max_size
+        ):
+            await self.rotate_files()
+
+        async with aiofiles.open(self.current_file, "a", encoding="utf-8") as f:
+            await f.write(json.dumps(data, ensure_ascii=False) + "\n")
+
+    async def rotate_files(self):
+        # Remove the oldest backup if it exists
+        oldest_backup = f"{self.base_filename}.{self.backup_count}"
+        if os.path.exists(oldest_backup):
+            os.remove(oldest_backup)
+
+        # Rotate existing backups
+        for i in range(self.backup_count - 1, 0, -1):
+            src = f"{self.base_filename}.{i}"
+            dst = f"{self.base_filename}.{i + 1}"
+            if os.path.exists(src):
+                os.rename(src, dst)
+
+        # Rename current file to backup
+        os.rename(self.current_file, f"{self.base_filename}.1")
+
+
+class AsyncKeywordTableSimpleRetriever(KeywordTableSimpleRetriever):
+    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, self._retrieve, query_bundle)
+
+
+class SampleableEmbeddingQADataset:
+    def __init__(self, dataset: EmbeddingQAFinetuneDataset):
+        self.dataset = dataset
+
+    def sample(self, n: int) -> EmbeddingQAFinetuneDataset:
+        """
+        Sample n queries from the dataset.
+
+        Args:
+            n (int): Number of queries to sample.
+
+        Returns:
+            EmbeddingQAFinetuneDataset: A new dataset with the sampled queries.
+        """
+        if n > len(self.dataset.queries):
+            raise ValueError(
+                f"n ({n}) is greater than the number of queries ({len(self.dataset.queries)})"
+            )
+
+        sampled_query_ids = random.sample(list(self.dataset.queries.keys()), n)
+
+        sampled_queries = {qid: self.dataset.queries[qid] for qid in sampled_query_ids}
+        sampled_relevant_docs = {
+            qid: self.dataset.relevant_docs[qid] for qid in sampled_query_ids
+        }
+
+        # Collect all unique document IDs from the sampled relevant docs
+        sampled_doc_ids = set()
+        for doc_ids in sampled_relevant_docs.values():
+            sampled_doc_ids.update(doc_ids)
+
+        sampled_corpus = {
+            doc_id: self.dataset.corpus[doc_id] for doc_id in sampled_doc_ids
+        }
+
+        return EmbeddingQAFinetuneDataset(
+            queries=sampled_queries,
+            corpus=sampled_corpus,
+            relevant_docs=sampled_relevant_docs,
+            mode=self.dataset.mode,
+        )
+
+    def __getattr__(self, name):
+        return getattr(self.dataset, name)
+
+
+class RetrieverEvaluator(BaseRetrievalEvaluator):
+    """Retriever evaluator.
+
+    This module will evaluate a retriever using a set of metrics.
+
+    Args:
+        metrics (List[BaseRetrievalMetric]): Sequence of metrics to evaluate
+        retriever: Retriever to evaluate.
+        node_postprocessors (Optional[List[BaseNodePostprocessor]]): Post-processor to apply after retrieval.
+    """
+
+    retriever: BaseRetriever = Field(..., description="Retriever to evaluate")
+    node_postprocessors: Optional[List[SerializeAsAny[BaseNodePostprocessor]]] = Field(
+        default=None, description="Optional post-processor"
+    )
+
+    async def _aget_retrieved_ids_and_texts(
+        self,
+        query: str,
+        mode: RetrievalEvalMode = RetrievalEvalMode.TEXT,
+        source: str = "",
+    ) -> Tuple[List[str], List[str]]:
+        """Get retrieved ids and texts, potentially applying a post-processor."""
+        try:
+            retrieved_nodes: list[NodeWithScore] = await self.retriever.aretrieve(query)
+            logfire.info(f"Retrieved {len(retrieved_nodes)} nodes for: '{query}'")
+        except Exception as e:
+            return ["00000000-0000-0000-0000-000000000000"], [str(e)]
+
+        if len(retrieved_nodes) == 0 or retrieved_nodes is None:
+            print(f"No nodes retrieved for {query}")
+            return ["00000000-0000-0000-0000-000000000000"], ["No nodes retrieved"]
+
+        if self.node_postprocessors:
+            for node_postprocessor in self.node_postprocessors:
+                retrieved_nodes = node_postprocessor.postprocess_nodes(
+                    retrieved_nodes, query_str=query
+                )
+
+        return (
+            [node.node.node_id for node in retrieved_nodes],
+            [node.node.text for node in retrieved_nodes],  # type: ignore
+        )
+
+
+class OpenAIAgentRetrieverEvaluator(BaseRetrievalEvaluator):
+    agent: OpenAIAgent = Field(description="The OpenAI agent used for retrieval")
+
+    async def _aget_retrieved_ids_and_texts(
+        self,
+        query: str,
+        mode: RetrievalEvalMode = RetrievalEvalMode.TEXT,
+        source: str = "",
+    ) -> Tuple[List[str], List[str]]:
+
+        self.agent.memory.reset()
+
+        try:
+            logfire.info(f"Executing agent with query: {query}")
+            response: AgentChatResponse = await self.agent.achat(query)
+        except Exception as e:
+            # await self._save_response_data_async(
+            #     source, query, ["Error retrieving nodes"], "Error retrieving nodes"
+            # )
+            return ["00000000-0000-0000-0000-000000000000"], [str(e)]
+
+        retrieved_nodes: list[NodeWithScore] = get_nodes_with_score(response)
+        logfire.info(f"Retrieved {len(retrieved_nodes)} to answer: '{query}'")
+        retrieved_nodes = retrieved_nodes[:6]  # Limit to first 6 retrieved nodes
+
+        if len(retrieved_nodes) == 0 or retrieved_nodes is None:
+            # await self._save_response_data_async(
+            #     source, query, ["No retrieved nodes"], "No retrieved nodes"
+            # )
+            return ["00000000-0000-0000-0000-000000000000"], ["No nodes retrieved"]
+
+        retrieved_ids = [node.node.node_id for node in retrieved_nodes]
+        retrieved_texts = [node.node.text for node in retrieved_nodes]  # type: ignore
+
+        # Will not save context as its too long (token wise), costly and takes too much time.
+        await self._save_response_data_async(
+            source=source, query=query, context="", response=response.response
+        )
+
+        return retrieved_ids, retrieved_texts
+
+    async def _save_response_data_async(self, source, query, context, response):
+        data = {
+            "source": source,
+            "question": query,
+            # "context": context,
+            "answer": response,
+        }
+        await rotating_writer.write(data)
+
+
+def get_nodes_with_score(completion) -> list[NodeWithScore]:
+    retrieved_nodes = []
+    for source in completion.sources:  # completion.sources = list[ToolOutput]
+        if source.is_error == True:
+            continue
+        for node in source.raw_output:  # source.raw_output = list[NodeWithScore]
+            retrieved_nodes.append(node)
+    return retrieved_nodes
+
+
+def setup_basic_database(db_collection, dict_file_name, keyword_retriever):
+    db = chromadb.PersistentClient(path=f"data/{db_collection}")
+    chroma_collection = db.get_or_create_collection(db_collection)
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+
+    # embed_model = OpenAIEmbedding(model="text-embedding-3-large", mode="similarity")
+    embed_model = CohereEmbedding(
+        api_key=os.environ["COHERE_API_KEY"],
+        model_name="embed-english-v3.0",
+        input_type="search_query",
+    )
+    # client = embed_model._get_client()
+    # aclient = embed_model._get_aclient()
+    # logfire.instrument_openai(client)
+    # logfire.instrument_openai(aclient)
+
+    index = VectorStoreIndex.from_vector_store(
+        vector_store=vector_store,
+        show_progress=True,
+    )
+    vector_retriever = VectorIndexRetriever(
+        index=index,
+        similarity_top_k=15,
+        embed_model=embed_model,
+    )
+    with open(f"data/{db_collection}/{dict_file_name}", "rb") as f:
+        document_dict = pickle.load(f)
+
+    return CustomRetriever(vector_retriever, document_dict, keyword_retriever, "OR")
+
+
+def update_query_engine_tools(selected_sources, custom_retriever_all_sources):
+    tools = []
+    source_mapping = {
+        # "Transformers Docs": (
+        #     custom_retriever_transformers,
+        #     "Transformers_information",
+        #     """Useful for general questions asking about the artificial intelligence (AI) field. Employ this tool to fetch information on topics such as language models (LLMs) models such as Llama3 and theory (transformer architectures), tips on prompting, quantization, etc.""",
+        # ),
+        # "PEFT Docs": (
+        #     custom_retriever_peft,
+        #     "PEFT_information",
+        #     """Useful for questions asking about efficient LLM fine-tuning. Employ this tool to fetch information on topics such as LoRA, QLoRA, etc.""",
+        # ),
+        # "TRL Docs": (
+        #     custom_retriever_trl,
+        #     "TRL_information",
+        #     """Useful for questions asking about fine-tuning LLMs with reinforcement learning (RLHF). Includes information about the Supervised Fine-tuning step (SFT), Reward Modeling step (RM), and the Proximal Policy Optimization (PPO) step.""",
+        # ),
+        # "LlamaIndex Docs": (
+        #     custom_retriever_llama_index,
+        #     "LlamaIndex_information",
+        #     """Useful for questions asking about retrieval augmented generation (RAG) with LLMs and embedding models. It is the documentation of a framework, includes info about fine-tuning embedding models, building chatbots, and agents with llms, using vector databases, embeddings, information retrieval with cosine similarity or bm25, etc.""",
+        # ),
+        # "OpenAI Cookbooks": (
+        #     custom_retriever_openai_cookbooks,
+        #     "openai_cookbooks_info",
+        #     """Useful for questions asking about accomplishing common tasks with the OpenAI API. Returns example code and guides stored in Jupyter notebooks, including info about ChatGPT GPT actions, OpenAI Assistants API,  and How to fine-tune OpenAI's GPT-4o and GPT-4o-mini models with the OpenAI API.""",
+        # ),
+        # "LangChain Docs": (
+        #     custom_retriever_langchain,
+        #     "langchain_info",
+        #     """Useful for questions asking about the LangChain framework. It is the documentation of the LangChain framework, includes info about building chains, agents, and tools, using memory, prompts, callbacks, etc.""",
+        # ),
+        "All Sources": (
+            custom_retriever_all_sources,
+            "all_sources_info",
+            """Useful for all questions, contains information about the field of AI.""",
+        ),
+    }
+
+    for source in selected_sources:
+        if source in source_mapping:
+            retriever, name, description = source_mapping[source]
+            tools.append(
+                RetrieverTool(
+                    retriever=retriever,
+                    metadata=ToolMetadata(
+                        name=name,
+                        description=description,
+                    ),
+                )
+            )
+
+    return tools
+
+
+def setup_agent(custom_retriever_all_sources) -> OpenAIAgent:
+
+    llm = OpenAI(
+        temperature=1,
+        # model="gpt-4o",
+        model="gpt-4o-mini",
+        max_tokens=5000,
+        max_retries=3,
+    )
+    client = llm._get_client()
+    logfire.instrument_openai(client)
+    aclient = llm._get_aclient()
+    logfire.instrument_openai(aclient)
+
+    tools_available = [
+        # "Transformers Docs",
+        # "PEFT Docs",
+        # "TRL Docs",
+        # "LlamaIndex Docs",
+        # "LangChain Docs",
+        # "OpenAI Cookbooks",
+        "All Sources",
+    ]
+    query_engine_tools = update_query_engine_tools(
+        tools_available, custom_retriever_all_sources
+    )
+
+    agent = OpenAIAgent.from_tools(
+        llm=llm,
+        tools=query_engine_tools,
+        system_prompt=system_message_openai_agent,
+    )
+
+    return agent
+
+
+async def evaluate_answers():
+    start_time = time.time()
+
+    # Gemini is not async here, maybe it could work with multithreading?
+    # llm = Gemini(model="models/gemini-1.5-flash-002", temperature=1, max_tokens=1000)
+    llm = OpenAI(model="gpt-4o-mini", temperature=1, max_tokens=1000)
+    relevancy_evaluator = AnswerRelevancyEvaluator(llm=llm)
+
+    # Load queries and response strings from JSONL file
+    query_response_pairs = []
+    with open("response_data.jsonl", "r") as f:
+        for line in f:
+            data = json.loads(line)
+            query_response_pairs.append(
+                (data["source"], data["query"], data["response"])
+            )
+
+    logfire.info(f"Number of queries and answers: {len(query_response_pairs)}")
+
+    semaphore = asyncio.Semaphore(90)  # Adjust this value as needed
+
+    async def evaluate_query_response(source, query, response):
+        async with semaphore:
+            try:
+                result: EvaluationResult = await relevancy_evaluator.aevaluate(
+                    query=query, response=response
+                )
+                return source, result
+            except Exception as e:
+                logfire.error(f"Error evaluating query for {source}: {str(e)}")
+                return source, None
+
+    # Use asyncio.gather to run all evaluations concurrently
+    results = await tqdm_asyncio.gather(
+        *[
+            evaluate_query_response(source, query, response)
+            for source, query, response in query_response_pairs
+        ],
+        desc="Evaluating answers",
+        total=len(query_response_pairs),
+    )
+
+    # Process results
+    eval_results = {}
+    for item in results:
+        if isinstance(item, tuple) and len(item) == 2:
+            source, result = item
+            if result is not None:
+                if source not in eval_results:
+                    eval_results[source] = []
+                eval_results[source].append(result)
+        else:
+            logfire.error(f"Unexpected result: {item}")
+
+    # Save results for each source
+    for source, results in eval_results.items():
+        with open(f"eval_answers_results_{source}.pkl", "wb") as f:
+            pickle.dump(results, f)
+
+    end_time = time.time()
+    logfire.info(f"Total evaluation time: {round(end_time - start_time, 3)} seconds")
+
+    return eval_results
+
+
+def create_docs(input_file: str) -> List[Document]:
+    with open(input_file, "r") as f:
+        documents = []
+        for line in f:
+            data = json.loads(line)
+            documents.append(
+                Document(
+                    doc_id=data["doc_id"],
+                    text=data["content"],
+                    metadata={  # type: ignore
+                        "url": data["url"],
+                        "title": data["name"],
+                        "tokens": data["tokens"],
+                        "retrieve_doc": data["retrieve_doc"],
+                        "source": data["source"],
+                    },
+                    excluded_llm_metadata_keys=[
+                        "title",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                    excluded_embed_metadata_keys=[
+                        "url",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                )
+            )
+    return documents
+
+
+def get_sample_size(source: str, total_queries: int) -> int:
+    """Determine the number of queries to sample based on the source."""
+    # small_datasets = {"peft": 0, "trl": 0, "openai_cookbooks": 0}
+    # large_datasets = {
+    #     "transformers": 0,
+    #     "llama_index": 0,
+    #     "langchain": 1,
+    #     "tai_blog": 0,
+    # }
+    small_datasets = {"peft": 49, "trl": 34, "openai_cookbooks": 170}
+    large_datasets = {
+        "transformers": 200,
+        "llama_index": 200,
+        "langchain": 200,
+        "tai_blog": 200,
+    }
+    # small_datasets = {"peft": 49, "trl": 34, "openai_cookbooks": 100}
+    # large_datasets = {
+    #     "transformers": 100,
+    #     "llama_index": 100,
+    #     "langchain": 100,
+    #     "tai_blog": 100,
+    # }
+    # small_datasets = {"peft": 18, "trl": 12, "openai_cookbooks": 14}
+    # large_datasets = {
+    #     "transformers": 24,
+    #     "llama_index": 8,
+    #     "langchain": 6,
+    #     "tai_blog": 18,
+    # }
+    # small_datasets = {"peft": 4, "trl": 4, "openai_cookbooks": 4}
+    # large_datasets = {
+    #     "transformers": 4,
+    #     "llama_index": 4,
+    #     "langchain": 5,
+    #     "tai_blog": 5,
+    # }
+
+    if source in small_datasets:
+        return small_datasets[source]
+    elif source in large_datasets:
+        return large_datasets[source]
+    else:
+        return min(100, total_queries)  # Default to 100 or all queries if less than 100
+
+
+async def evaluate_retriever():
+    start_time = time.time()
+    with open("data/keyword_retriever_async.pkl", "rb") as f:
+        keyword_retriever = pickle.load(f)
+
+    custom_retriever_all_sources: CustomRetriever = setup_basic_database(
+        "chroma-db-all_sources", "document_dict_all_sources.pkl", keyword_retriever
+    )
+    # agent = setup_agent(custom_retriever_all_sources)
+
+    # filters = MetadataFilters(
+    #     filters=[
+    #         MetadataFilter(key="source", operator=FilterOperator.EQ, value="langchain"),
+    #     ]
+    # )
+    # custom_retriever_all_sources._vector_retriever._filters = filters
+
+    end_time = time.time()
+    logfire.info(
+        f"Time taken for setup the custom retriever: {round(end_time - start_time, 2)} seconds"
+    )
+
+    sources_to_evaluate = [
+        "transformers",
+        "peft",
+        "trl",
+        "llama_index",
+        "langchain",
+        "openai_cookbooks",
+        "tai_blog",
+    ]
+
+    # for k in [5, 7, 9, 11, 13, 15]:
+    #     custom_retriever_all_sources._vector_retriever._similarity_top_k = k
+
+    retriever_evaluator = RetrieverEvaluator.from_metric_names(
+        ["mrr", "hit_rate"], retriever=custom_retriever_all_sources
+    )
+    # retriever_evaluator = OpenAIAgentRetrieverEvaluator.from_metric_names(
+    #     metric_names=["mrr", "hit_rate"], agent=agent
+    # )
+
+    all_query_pairs = []
+    for source in sources_to_evaluate:
+        rag_eval_dataset = EmbeddingQAFinetuneDataset.from_json(
+            f"scripts/rag_eval_{source}.json"
+        )
+        sampleable_dataset = SampleableEmbeddingQADataset(rag_eval_dataset)
+        sample_size = get_sample_size(source, len(sampleable_dataset.queries))
+        sampled_dataset = sampleable_dataset.sample(n=sample_size)
+        query_expected_ids_pairs = sampled_dataset.query_docid_pairs
+        all_query_pairs.extend(
+            [(source, pair[0], pair[1]) for pair in query_expected_ids_pairs]
+        )
+
+    semaphore = asyncio.Semaphore(220)  # 250 caused a couple of errors
+    # semaphore = asyncio.Semaphore(90)  # 100 caused a couple of errors with agent
+
+    async def evaluate_query(source, query, expected_ids):
+        async with semaphore:
+            try:
+                result: RetrievalEvalResult = await retriever_evaluator.aevaluate(
+                    query=query,
+                    expected_ids=expected_ids,
+                    mode=RetrievalEvalMode.TEXT,
+                    source=source,
+                )
+                return source, result
+            except Exception as e:
+                logfire.error(f"Error evaluating query for {source}: {str(e)}")
+                return source, None
+
+    # Use asyncio.gather to run all evaluations concurrently
+    results = await tqdm_asyncio.gather(
+        *[
+            evaluate_query(source, query, expected_ids)
+            for source, query, expected_ids in all_query_pairs
+        ],
+        desc="Evaluating queries",
+        total=len(all_query_pairs),
+    )
+
+    # Process results
+    eval_results = {source: [] for source in sources_to_evaluate}
+    for item in results:
+        if isinstance(item, tuple) and len(item) == 2:
+            source, result = item
+            if result is not None:
+                eval_results[source].append(result)
+        else:
+            logfire.error(f"Unexpected result: {item}")
+
+    # Save results for each source
+    for source, results in eval_results.items():
+        with open(f"eval_results_{source}.pkl", "wb") as f:
+            pickle.dump(results, f)
+        # print(display_results_retriever(source, results))
+
+    end_time = time.time()
+    logfire.info(f"Total evaluation time: {round(end_time - start_time, 3)} seconds")
+
+
+def display_results_retriever(name, eval_results):
+    """Display results from evaluate."""
+
+    metric_dicts = []
+    for eval_result in eval_results:
+        metric_dict = eval_result.metric_vals_dict
+        metric_dicts.append(metric_dict)
+
+    full_df = pd.DataFrame(metric_dicts)
+
+    hit_rate = full_df["hit_rate"].mean()
+    mrr = full_df["mrr"].mean()
+
+    metric_df = pd.DataFrame(
+        {"Retriever Name": [name], "Hit Rate": [hit_rate], "MRR": [mrr]}
+    )
+
+    return metric_df
+
+
+def display_results():
+
+    sources = [
+        "transformers",
+        "peft",
+        "trl",
+        "llama_index",
+        "langchain",
+        "openai_cookbooks",
+        "tai_blog",
+    ]
+    # retrievers_to_evaluate = [
+    #     # "chroma-db-all_sources_400_0",
+    #     # "chroma-db-all_sources_400_200",
+    #     # "chroma-db-all_sources_500_0",
+    #     # "chroma-db-all_sources_500_250",
+    #     # "chroma-db-all_sources",
+    #     # "chroma-db-all_sources_800_400",
+    #     # "chroma-db-all_sources_1000_0",
+    #     # "chroma-db-all_sources_1000_500",
+    # ]
+
+    # topk = [5, 7, 9, 11, 13, 15]
+    # for k in topk:
+    # for db_name in retrievers_to_evaluate:
+    if True:
+        # print("-" * 20)
+        # print(f"Retriever {db_name}")
+        for source in sources:
+            with open(f"eval_results_{source}.pkl", "rb") as f:
+                eval_results = pickle.load(f)
+            print(display_results_retriever(f"{source}", eval_results))
+
+
+def display_results_answers():
+
+    sources = [
+        "transformers",
+        "peft",
+        "trl",
+        "llama_index",
+        "langchain",
+        "openai_cookbooks",
+        "tai_blog",
+    ]
+
+    for source in sources:
+        with open(f"eval_answers_results_{source}.pkl", "rb") as f:
+            eval_results = pickle.load(f)
+        print(
+            f"Score for {source}:",
+            sum(result.score for result in eval_results) / len(eval_results),
+        )
+
+
+async def main():
+    await evaluate_retriever()
+    display_results()
+    # await evaluate_answers()
+    # display_results_answers()
+    return
+
+
+if __name__ == "__main__":
+
+    logfire.configure()
+    rotating_writer = RotatingJSONLWriter(
+        "response_data.jsonl", max_size=10**7, backup_count=5
+    )
+
+    start_time = time.time()
+    asyncio.run(main())
+    end_time = time.time()
+    logfire.info(
+        f"Time taken to run script: {round((end_time - start_time), 3)} seconds"
+    )
+
+    # # Creating the keyword index and retriever
+    # logfire.info("Creating nodes from documents")
+    # documents = create_docs("data/all_sources_data.jsonl")
+    # pipeline = IngestionPipeline(
+    #     transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)]
+    # )
+    # all_nodes = pipeline.run(documents=documents, show_progress=True)
+    # # with open("data/all_nodes.pkl", "wb") as f:
+    # #     pickle.dump(all_nodes, f)
+
+    # # all_nodes = pickle.load(open("data/all_nodes.pkl", "rb"))
+    # logfire.info(f"Number of nodes: {len(all_nodes)}")
+
+    # with open("processed_chunks.pkl", "rb") as f:
+    #     all_nodes: list[TextNode] = pickle.load(f)
+
+    # keyword_index = SimpleKeywordTableIndex(
+    #     nodes=all_nodes, max_keywords_per_chunk=10, show_progress=True, use_async=False
+    # )
+    # # with open("data/keyword_index.pkl", "wb") as f:
+    # #     pickle.dump(keyword_index, f)
+    # # keyword_index = pickle.load(open("data/keyword_index.pkl", "rb"))
+
+    # logfire.info("Creating keyword retriever")
+    # keyword_retriever = AsyncKeywordTableSimpleRetriever(index=keyword_index)
+
+    # with open("data/keyword_retriever_async.pkl", "wb") as f:
+    #     pickle.dump(keyword_retriever, f)

scripts/generate_qa_dataset.ipynb

@@ -0,0 +1,287 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import nest_asyncio\n",
+    "\n",
+    "nest_asyncio.apply()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Generate synthetic dataset of questions + doc ids\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# DEFAULT_QA_GENERATE_PROMPT_TMPL = \"\"\"\\\n",
+    "# Context information is below.\n",
+    "\n",
+    "# ---------------------\n",
+    "# {context_str}\n",
+    "# ---------------------\n",
+    "\n",
+    "# Given the context information and not prior knowledge.\n",
+    "# generate only questions based on the below query.\n",
+    "\n",
+    "# You are a Teacher/ Professor with expertise in the field of AI. Your task is to setup \\\n",
+    "# {num_questions_per_chunk} questions for an upcoming \\\n",
+    "# quiz/examination. The questions should be diverse in nature \\\n",
+    "# across the document. Restrict the questions to the \\\n",
+    "# context information provided.\"\n",
+    "# \"\"\"\n",
+    "\n",
+    "STUDENT_AI_QUESTIONS_GENERATE_PROMPT_TMPL = \"\"\"\\\n",
+    "Context information is below.\n",
+    "\n",
+    "---------------------\n",
+    "{context_str}\n",
+    "---------------------\n",
+    "\n",
+    "Given the context information above, generate {num_questions_per_chunk} questions that a student might ask about AI, specifically related to the information provided in the context, \n",
+    "but without the questions mentioning the context information.\n",
+    "\n",
+    "You are simulating curious students who are learning about AI. Your task is to create questions that:\n",
+    "1. Reflect genuine curiosity about the topics covered in the context.\n",
+    "2. Vary in complexity, from basic clarifications to more advanced inquiries.\n",
+    "3. Demonstrate a student's desire to understand AI concepts better.\n",
+    "4. Can be answered using the information provided in the context.\n",
+    "\n",
+    "The questions should be diverse and cover different aspects of the content. Do not use any prior knowledge beyond what's given in the context. Ensure that each question you generate can be answered using the information provided in the context.\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llama_index.llms.openai import OpenAI\n",
+    "from llama_index.llms.gemini import Gemini\n",
+    "\n",
+    "# llm = Gemini(model=\"models/gemini-1.5-flash-latest\", temperature=1, max_tokens=1000)\n",
+    "llm = OpenAI(\n",
+    "    api_key=\"\",\n",
+    "    model=\"gpt-4o-mini\",\n",
+    "    temperature=1,\n",
+    "    max_tokens=1000,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import re\n",
+    "import uuid\n",
+    "import warnings\n",
+    "import asyncio\n",
+    "from typing import Dict, List, Tuple\n",
+    "\n",
+    "from llama_index.core.llms.llm import LLM\n",
+    "from llama_index.core.schema import MetadataMode, TextNode\n",
+    "from tqdm.asyncio import tqdm as async_tqdm\n",
+    "from llama_index.core.llama_dataset.legacy.embedding import EmbeddingQAFinetuneDataset\n",
+    "from llama_index.core import Document\n",
+    "\n",
+    "\n",
+    "async def generate_qa_embedding_pairs(\n",
+    "    nodes: List[TextNode],\n",
+    "    llm: LLM,\n",
+    "    qa_generate_prompt_tmpl: str = STUDENT_AI_QUESTIONS_GENERATE_PROMPT_TMPL,\n",
+    "    num_questions_per_chunk: int = 1,\n",
+    "    max_concurrent: int = 10,\n",
+    "    delay: float = 0.5,\n",
+    ") -> EmbeddingQAFinetuneDataset:\n",
+    "    \"\"\"Generate examples given a set of nodes.\"\"\"\n",
+    "    node_dict = {\n",
+    "        node.node_id: node.get_content(metadata_mode=MetadataMode.NONE)\n",
+    "        for node in nodes\n",
+    "    }\n",
+    "\n",
+    "    queries = {}\n",
+    "    relevant_docs = {}\n",
+    "    semaphore = asyncio.Semaphore(max_concurrent)\n",
+    "\n",
+    "    async def process_node(node_id: str, text: str):\n",
+    "        async with semaphore:\n",
+    "            query = qa_generate_prompt_tmpl.format(\n",
+    "                context_str=text, num_questions_per_chunk=num_questions_per_chunk\n",
+    "            )\n",
+    "            response = await llm.acomplete(\n",
+    "                query\n",
+    "            )  # Assuming the LLM has an async method\n",
+    "\n",
+    "            result = str(response).strip().split(\"\\n\")\n",
+    "            questions = [\n",
+    "                re.sub(r\"^\\d+[\\).\\s]\", \"\", question).strip() for question in result\n",
+    "            ]\n",
+    "            questions = [question for question in questions if len(question) > 0][\n",
+    "                :num_questions_per_chunk\n",
+    "            ]\n",
+    "\n",
+    "            num_questions_generated = len(questions)\n",
+    "            if num_questions_generated < num_questions_per_chunk:\n",
+    "                warnings.warn(\n",
+    "                    f\"Fewer questions generated ({num_questions_generated}) \"\n",
+    "                    f\"than requested ({num_questions_per_chunk}).\"\n",
+    "                )\n",
+    "\n",
+    "            for question in questions:\n",
+    "                question_id = str(uuid.uuid4())\n",
+    "                queries[question_id] = question\n",
+    "                relevant_docs[question_id] = [node_id]\n",
+    "\n",
+    "            await asyncio.sleep(delay)\n",
+    "\n",
+    "    # Use asyncio.gather to process nodes concurrently\n",
+    "    await async_tqdm.gather(\n",
+    "        *[process_node(node_id, text) for node_id, text in node_dict.items()]\n",
+    "    )\n",
+    "\n",
+    "    # construct dataset\n",
+    "    return EmbeddingQAFinetuneDataset(\n",
+    "        queries=queries, corpus=node_dict, relevant_docs=relevant_docs\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pickle\n",
+    "\n",
+    "\n",
+    "async def generate_questions(path, source_name):\n",
+    "    nodes = []\n",
+    "    with open(path, \"rb\") as f:\n",
+    "        document_dict = pickle.load(f)\n",
+    "        for doc_id in document_dict.keys():\n",
+    "            doc: Document = document_dict[doc_id]\n",
+    "            if doc.metadata[\"tokens\"] >= 100_000:\n",
+    "                print(\"skipping\", doc.metadata[\"tokens\"])\n",
+    "                continue\n",
+    "            node = TextNode(text=doc.text, metadata=doc.metadata, id_=doc_id)\n",
+    "            nodes.append(node)\n",
+    "\n",
+    "    rag_eval_dataset: EmbeddingQAFinetuneDataset = await generate_qa_embedding_pairs(\n",
+    "        nodes,\n",
+    "        llm=llm,\n",
+    "        num_questions_per_chunk=1,\n",
+    "        qa_generate_prompt_tmpl=STUDENT_AI_QUESTIONS_GENERATE_PROMPT_TMPL,\n",
+    "        max_concurrent=20,  # Adjust this to control concurrency\n",
+    "        delay=0.5,  # Adjust this to add delay between API calls\n",
+    "    )\n",
+    "    rag_eval_dataset.save_json(f\"./rag_eval_{source_name}.json\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# await generate_questions(\n",
+    "#     \"../data/chroma-db-langchain/document_dict_langchain.pkl\", \"langchain\"\n",
+    "# )\n",
+    "# await generate_questions(\n",
+    "#     \"../data/chroma-db-llama_index/document_dict_llama_index.pkl\", \"llama_index\"\n",
+    "# )\n",
+    "# await generate_questions(\n",
+    "#     \"../data/chroma-db-openai_cookbooks/document_dict_openai_cookbooks.pkl\",\n",
+    "#     \"openai_cookbooks\",\n",
+    "# )\n",
+    "# await generate_questions(\"../data/chroma-db-peft/document_dict_peft.pkl\", \"peft\")\n",
+    "# await generate_questions(\"../data/chroma-db-trl/document_dict_trl.pkl\", \"trl\")\n",
+    "\n",
+    "await generate_questions(\n",
+    "    \"../data/chroma-db-tai_blog/document_dict_tai_blog.pkl\", \"tai_blog\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # We can also load the dataset from a previously saved json file.\n",
+    "# from llama_index.core.evaluation import EmbeddingQAFinetuneDataset\n",
+    "\n",
+    "# rag_eval_dataset = EmbeddingQAFinetuneDataset.from_json(\"./rag_eval_transformers.json\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.12.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

scripts/main.py

@@ -0,0 +1,291 @@
+import pdb
+
+import gradio as gr
+import logfire
+from custom_retriever import CustomRetriever
+from llama_index.agent.openai import OpenAIAgent
+from llama_index.core.llms import MessageRole
+from llama_index.core.memory import ChatSummaryMemoryBuffer
+from llama_index.core.tools import RetrieverTool, ToolMetadata
+from llama_index.core.vector_stores import (
+    FilterCondition,
+    FilterOperator,
+    MetadataFilter,
+    MetadataFilters,
+)
+from llama_index.llms.openai import OpenAI
+from prompts import system_message_openai_agent
+from setup import (  # custom_retriever_langchain,; custom_retriever_llama_index,; custom_retriever_openai_cookbooks,; custom_retriever_peft,; custom_retriever_transformers,; custom_retriever_trl,
+    AVAILABLE_SOURCES,
+    AVAILABLE_SOURCES_UI,
+    CONCURRENCY_COUNT,
+    custom_retriever_all_sources,
+)
+
+
+def update_query_engine_tools(selected_sources) -> list[RetrieverTool]:
+    tools = []
+    source_mapping: dict[str, tuple[CustomRetriever, str, str]] = {
+        # "Transformers Docs": (
+        #     custom_retriever_transformers,
+        #     "Transformers_information",
+        #     """Useful for general questions asking about the artificial intelligence (AI) field. Employ this tool to fetch information on topics such as language models (LLMs) models such as Llama3 and theory (transformer architectures), tips on prompting, quantization, etc.""",
+        # ),
+        # "PEFT Docs": (
+        #     custom_retriever_peft,
+        #     "PEFT_information",
+        #     """Useful for questions asking about efficient LLM fine-tuning. Employ this tool to fetch information on topics such as LoRA, QLoRA, etc.""",
+        # ),
+        # "TRL Docs": (
+        #     custom_retriever_trl,
+        #     "TRL_information",
+        #     """Useful for questions asking about fine-tuning LLMs with reinforcement learning (RLHF). Includes information about the Supervised Fine-tuning step (SFT), Reward Modeling step (RM), and the Proximal Policy Optimization (PPO) step.""",
+        # ),
+        # "LlamaIndex Docs": (
+        #     custom_retriever_llama_index,
+        #     "LlamaIndex_information",
+        #     """Useful for questions asking about retrieval augmented generation (RAG) with LLMs and embedding models. It is the documentation of a framework, includes info about fine-tuning embedding models, building chatbots, and agents with llms, using vector databases, embeddings, information retrieval with cosine similarity or bm25, etc.""",
+        # ),
+        # "OpenAI Cookbooks": (
+        #     custom_retriever_openai_cookbooks,
+        #     "openai_cookbooks_info",
+        #     """Useful for questions asking about accomplishing common tasks with the OpenAI API. Returns example code and guides stored in Jupyter notebooks, including info about ChatGPT GPT actions, OpenAI Assistants API,  and How to fine-tune OpenAI's GPT-4o and GPT-4o-mini models with the OpenAI API.""",
+        # ),
+        # "LangChain Docs": (
+        #     custom_retriever_langchain,
+        #     "langchain_info",
+        #     """Useful for questions asking about the LangChain framework. It is the documentation of the LangChain framework, includes info about building chains, agents, and tools, using memory, prompts, callbacks, etc.""",
+        # ),
+        "All Sources": (
+            custom_retriever_all_sources,
+            "all_sources_info",
+            """Useful for questions asking about information in the field of AI.""",
+        ),
+    }
+
+    for source in selected_sources:
+        if source in source_mapping:
+            custom_retriever, name, description = source_mapping[source]
+            tools.append(
+                RetrieverTool(
+                    retriever=custom_retriever,
+                    metadata=ToolMetadata(
+                        name=name,
+                        description=description,
+                    ),
+                )
+            )
+
+    return tools
+
+
+def generate_completion(
+    query,
+    history,
+    sources,
+    model,
+    memory,
+):
+    llm = OpenAI(temperature=1, model=model, max_tokens=None)
+    client = llm._get_client()
+    logfire.instrument_openai(client)
+
+    with logfire.span(f"Running query: {query}"):
+        logfire.info(f"User chosen sources: {sources}")
+
+        memory_chat_list = memory.get()
+
+        if len(memory_chat_list) != 0:
+            user_index_memory = [
+                i
+                for i, msg in enumerate(memory_chat_list)
+                if msg.role == MessageRole.USER
+            ]
+
+            user_index_history = [
+                i for i, msg in enumerate(history) if msg["role"] == "user"
+            ]
+
+            if len(user_index_memory) > len(user_index_history):
+                logfire.warn(f"There are more user messages in memory than in history")
+                user_index_to_remove = user_index_memory[len(user_index_history)]
+                memory_chat_list = memory_chat_list[:user_index_to_remove]
+                memory.set(memory_chat_list)
+
+        logfire.info(f"chat_history: {len(memory.get())} {memory.get()}")
+        logfire.info(f"gradio_history: {len(history)} {history}")
+
+        query_engine_tools: list[RetrieverTool] = update_query_engine_tools(
+            ["All Sources"]
+        )
+
+        filter_list = []
+        source_mapping = {
+            "Transformers Docs": "transformers",
+            "PEFT Docs": "peft",
+            "TRL Docs": "trl",
+            "LlamaIndex Docs": "llama_index",
+            "LangChain Docs": "langchain",
+            "OpenAI Cookbooks": "openai_cookbooks",
+            "Towards AI Blog": "tai_blog",
+        }
+
+        for source in sources:
+            if source in source_mapping:
+                filter_list.append(
+                    MetadataFilter(
+                        key="source",
+                        operator=FilterOperator.EQ,
+                        value=source_mapping[source],
+                    )
+                )
+
+        filters = MetadataFilters(
+            filters=filter_list,
+            condition=FilterCondition.OR,
+        )
+        # logfire.info(f"Filters: {filters}")
+        query_engine_tools[0].retriever._vector_retriever._filters = filters
+
+        # pdb.set_trace()
+
+        agent = OpenAIAgent.from_tools(
+            llm=llm,
+            memory=memory,
+            tools=query_engine_tools,
+            system_prompt=system_message_openai_agent,
+        )
+
+        completion = agent.stream_chat(query)
+
+    answer_str = ""
+    for token in completion.response_gen:
+        answer_str += token
+        yield answer_str
+
+    for answer_str in add_sources(answer_str, completion):
+        yield answer_str
+
+
+def add_sources(answer_str, completion):
+    if completion is None:
+        yield answer_str
+
+    formatted_sources = format_sources(completion)
+    if formatted_sources == "":
+        yield answer_str
+
+    if formatted_sources != "":
+        answer_str += "\n\n" + formatted_sources
+
+    yield answer_str
+
+
+def format_sources(completion) -> str:
+    if len(completion.sources) == 0:
+        return ""
+
+    # logfire.info(f"Formatting sources: {completion.sources}")
+
+    display_source_to_ui = {
+        src: ui for src, ui in zip(AVAILABLE_SOURCES, AVAILABLE_SOURCES_UI)
+    }
+
+    documents_answer_template: str = (
+        "📝 Here are the sources I used to answer your question:\n{documents}"
+    )
+    document_template: str = "[🔗 {source}: {title}]({url}), relevance: {score:2.2f}"
+    all_documents = []
+    for source in completion.sources:  # looping over list[ToolOutput]
+        if isinstance(source.raw_output, Exception):
+            logfire.error(f"Error in source output: {source.raw_output}")
+            # pdb.set_trace()
+            continue
+
+        if not isinstance(source.raw_output, list):
+            logfire.warn(f"Unexpected source output type: {type(source.raw_output)}")
+            continue
+        for src in source.raw_output:  # looping over list[NodeWithScore]
+            document = document_template.format(
+                title=src.metadata["title"],
+                score=src.score,
+                source=display_source_to_ui.get(
+                    src.metadata["source"], src.metadata["source"]
+                ),
+                url=src.metadata["url"],
+            )
+            all_documents.append(document)
+
+    if len(all_documents) == 0:
+        return ""
+    else:
+        documents = "\n".join(all_documents)
+        return documents_answer_template.format(documents=documents)
+
+
+def save_completion(completion, history):
+    pass
+
+
+def vote(data: gr.LikeData):
+    pass
+
+
+accordion = gr.Accordion(label="Customize Sources (Click to expand)", open=False)
+sources = gr.CheckboxGroup(
+    AVAILABLE_SOURCES_UI,
+    label="Sources",
+    value=[
+        "Transformers Docs",
+        "PEFT Docs",
+        "TRL Docs",
+        "LlamaIndex Docs",
+        "LangChain Docs",
+        "OpenAI Cookbooks",
+        "Towards AI Blog",
+        # "All Sources",
+    ],
+    interactive=True,
+)
+model = gr.Dropdown(
+    [
+        "gpt-4o-mini",
+    ],
+    label="Model",
+    value="gpt-4o-mini",
+    interactive=False,
+)
+
+with gr.Blocks(
+    title="Towards AI 🤖",
+    analytics_enabled=True,
+    fill_height=True,
+) as demo:
+
+    memory = gr.State(
+        lambda: ChatSummaryMemoryBuffer.from_defaults(
+            token_limit=120000,
+        )
+    )
+    chatbot = gr.Chatbot(
+        type="messages",
+        scale=8,
+        placeholder="<strong>Towards AI 🤖: A Question-Answering Bot for anything AI-related</strong><br>",
+        show_label=False,
+        show_copy_button=True,
+    )
+    chatbot.like(vote, None, None)
+    gr.ChatInterface(
+        fn=generate_completion,
+        type="messages",
+        chatbot=chatbot,
+        additional_inputs=[sources, model, memory],
+        additional_inputs_accordion=accordion,
+        fill_height=True,
+        fill_width=True,
+        analytics_enabled=True,
+    )
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=CONCURRENCY_COUNT)
+    demo.launch(debug=False, share=False)

scripts/prompts.py

@@ -0,0 +1,29 @@
+# Prompt 5
+system_message_openai_agent = """You are an AI teacher, answering questions from students of an applied AI course on Large Language Models (LLMs or llm) and Retrieval Augmented Generation (RAG) for LLMs.
+
+Topics covered include training models, fine-tuning models, giving memory to LLMs, prompting tips, hallucinations and bias, vector databases, transformer architectures, embeddings, RAG frameworks such as Langchain and LlamaIndex, making LLMs interact with tools, AI agents, reinforcement learning with human feedback (RLHF). Questions should be understood in this context.
+
+Your answers are aimed to teach students, so they should be complete, clear, and easy to understand.
+
+Use the available tools to gather insights pertinent to the field of AI.
+
+To answer student questions, always use the all_sources_info tool.
+
+Only some information returned by the tools might be relevant to the question, so ignore the irrelevant part and answer the question with what you have.
+
+Your responses are exclusively based on the output provided by the tools. Refrain from incorporating information not directly obtained from the tool's responses.
+
+When the conversation deepens or shifts focus within a topic, adapt your input to the tools to reflect these nuances. This means if a user requests further elaboration on a specific aspect of a previously discussed topic, you should reformulate your input to the tool to capture this new angle or more profound layer of inquiry.
+
+Provide comprehensive answers, ideally structured in multiple paragraphs, drawing from the tool's variety of relevant details. The depth and breadth of your responses should align with the scope and specificity of the information retrieved.
+
+Should the tools repository lack information on the queried topic, politely inform the user that the question transcends the bounds of your current knowledge base, citing the absence of relevant content in the tool's documentation.
+
+At the end of your answers, always invite the students to ask deeper questions about the topic if they have any. Make sure reformulate the question to the tool to capture this new angle or more profound layer of inquiry.
+
+Do not refer to the documentation directly, but use the information provided within it to answer questions.
+
+If code is provided in the information, share it with the students. It's important to provide complete code blocks so they can execute the code when they copy and paste them.
+
+Make sure to format your answers in Markdown format, including code blocks and snippets.
+"""

scripts/setup.py

@@ -0,0 +1,212 @@
+import asyncio
+import json
+import logging
+import os
+import pickle
+
+import chromadb
+import logfire
+from custom_retriever import CustomRetriever
+from dotenv import load_dotenv
+from llama_index.core import Document, SimpleKeywordTableIndex, VectorStoreIndex
+from llama_index.core.ingestion import IngestionPipeline
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.core.retrievers import (
+    KeywordTableSimpleRetriever,
+    VectorIndexRetriever,
+)
+from llama_index.core.schema import NodeWithScore, QueryBundle
+from llama_index.embeddings.cohere import CohereEmbedding
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from utils import init_mongo_db
+
+load_dotenv()
+
+logfire.configure()
+
+if not os.path.exists("data/chroma-db-all_sources"):
+    # Download the vector database from the Hugging Face Hub if it doesn't exist locally
+    # https://huggingface.co/datasets/towardsai-buster/ai-tutor-vector-db/tree/main
+    logfire.warn(
+        f"Vector database does not exist at 'data/chroma-db-all_sources', downloading from Hugging Face Hub"
+    )
+    from huggingface_hub import snapshot_download
+
+    snapshot_download(
+        repo_id="towardsai-buster/ai-tutor-vector-db",
+        local_dir="data",
+        repo_type="dataset",
+    )
+    logfire.info(f"Downloaded vector database to 'data/chroma-db-all_sources'")
+
+
+def create_docs(input_file: str) -> list[Document]:
+    with open(input_file, "r") as f:
+        documents = []
+        for line in f:
+            data = json.loads(line)
+            documents.append(
+                Document(
+                    doc_id=data["doc_id"],
+                    text=data["content"],
+                    metadata={  # type: ignore
+                        "url": data["url"],
+                        "title": data["name"],
+                        "tokens": data["tokens"],
+                        "retrieve_doc": data["retrieve_doc"],
+                        "source": data["source"],
+                    },
+                    excluded_llm_metadata_keys=[
+                        "title",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                    excluded_embed_metadata_keys=[
+                        "url",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                )
+            )
+    return documents
+
+
+def setup_database(db_collection, dict_file_name) -> CustomRetriever:
+    db = chromadb.PersistentClient(path=f"data/{db_collection}")
+    chroma_collection = db.get_or_create_collection(db_collection)
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    embed_model = CohereEmbedding(
+        api_key=os.environ["COHERE_API_KEY"],
+        model_name="embed-english-v3.0",
+        input_type="search_query",
+    )
+
+    index = VectorStoreIndex.from_vector_store(
+        vector_store=vector_store,
+        transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)],
+        show_progress=True,
+        # use_async=True,
+    )
+    vector_retriever = VectorIndexRetriever(
+        index=index,
+        similarity_top_k=15,
+        embed_model=embed_model,
+        # use_async=True,
+    )
+    with open(f"data/{db_collection}/{dict_file_name}", "rb") as f:
+        document_dict = pickle.load(f)
+
+    # with open("data/keyword_retriever_sync.pkl", "rb") as f:
+    # keyword_retriever: KeywordTableSimpleRetriever = pickle.load(f)
+
+    # keyword_retriever.num_chunks_per_query = 15
+
+    # # Creating the keyword index and retriever
+    # logfire.info("Creating nodes from documents")
+    # documents = create_docs("data/all_sources_data.jsonl")
+    # pipeline = IngestionPipeline(
+    #     transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)]
+    # )
+    # all_nodes = pipeline.run(documents=documents, show_progress=True)
+    # # with open("data/all_nodes.pkl", "wb") as f:
+    # #     pickle.dump(all_nodes, f)
+
+    # all_nodes = pickle.load(open("data/nodes_with_added_context.pkl", "rb"))
+    # logfire.info(f"Number of nodes: {len(all_nodes)}")
+
+    # keyword_index = SimpleKeywordTableIndex(
+    #     nodes=all_nodes, max_keywords_per_chunk=10, show_progress=True, use_async=False
+    # )
+    # # with open("data/keyword_index.pkl", "wb") as f:
+    # # pickle.dump(keyword_index, f)
+
+    # # keyword_index = pickle.load(open("data/keyword_index.pkl", "rb"))
+
+    # logfire.info("Creating keyword retriever")
+    # keyword_retriever = KeywordTableSimpleRetriever(index=keyword_index)
+
+    # with open("data/keyword_retriever_sync.pkl", "wb") as f:
+    #     pickle.dump(keyword_retriever, f)
+
+    # 'OR' Means both the vector nodes and the keyword nodes
+    # return CustomRetriever(vector_retriever, document_dict, keyword_retriever, "OR")
+    return CustomRetriever(vector_retriever, document_dict)
+
+
+# Setup retrievers
+# custom_retriever_transformers: CustomRetriever = setup_database(
+#     "chroma-db-transformers",
+#     "document_dict_transformers.pkl",
+# )
+# custom_retriever_peft: CustomRetriever = setup_database(
+#     "chroma-db-peft", "document_dict_peft.pkl"
+# )
+# custom_retriever_trl: CustomRetriever = setup_database(
+#     "chroma-db-trl", "document_dict_trl.pkl"
+# )
+# custom_retriever_llama_index: CustomRetriever = setup_database(
+#     "chroma-db-llama_index",
+#     "document_dict_llama_index.pkl",
+# )
+# custom_retriever_openai_cookbooks: CustomRetriever = setup_database(
+#     "chroma-db-openai_cookbooks",
+#     "document_dict_openai_cookbooks.pkl",
+# )
+# custom_retriever_langchain: CustomRetriever = setup_database(
+#     "chroma-db-langchain",
+#     "document_dict_langchain.pkl",
+# )
+
+custom_retriever_all_sources: CustomRetriever = setup_database(
+    "chroma-db-all_sources",
+    "document_dict_all_sources.pkl",
+)
+
+# Constants
+CONCURRENCY_COUNT = int(os.getenv("CONCURRENCY_COUNT", 64))
+MONGODB_URI = os.getenv("MONGODB_URI")
+
+AVAILABLE_SOURCES_UI = [
+    "Transformers Docs",
+    "PEFT Docs",
+    "TRL Docs",
+    "LlamaIndex Docs",
+    "LangChain Docs",
+    "OpenAI Cookbooks",
+    "Towards AI Blog",
+    # "All Sources",
+]
+
+AVAILABLE_SOURCES = [
+    "transformers",
+    "peft",
+    "trl",
+    "llama_index",
+    "langchain",
+    "openai_cookbooks",
+    "tai_blog",
+    # "all_sources",
+]
+
+mongo_db = (
+    init_mongo_db(uri=MONGODB_URI, db_name="towardsai-buster")
+    if MONGODB_URI
+    else logfire.warn("No mongodb uri found, you will not be able to save data.")
+)
+
+__all__ = [
+    # "custom_retriever_transformers",
+    # "custom_retriever_peft",
+    # "custom_retriever_trl",
+    # "custom_retriever_llama_index",
+    # "custom_retriever_openai_cookbooks",
+    # "custom_retriever_langchain",
+    "custom_retriever_all_sources",
+    "mongo_db",
+    "CONCURRENCY_COUNT",
+    "AVAILABLE_SOURCES_UI",
+    "AVAILABLE_SOURCES",
+]

scripts/utils.py

@@ -0,0 +1,16 @@
+from pymongo.mongo_client import MongoClient
+from pymongo.server_api import ServerApi
+
+
+def init_mongo_db(uri: str, db_name: str):
+    """Initialize the mongodb database."""
+
+    try:
+        assert uri is not None, "No URI passed"
+        client = MongoClient(uri, server_api=ServerApi("1"))
+        database = client[db_name]
+        print("Connected to MongoDB")
+        return database
+    except Exception as e:
+        print("Something went wrong connecting to mongodb")
+        return