Code to add metadata to the chunks

.chainlit/config.toml

@@ -22,7 +22,7 @@ prompt_playground = true
 unsafe_allow_html = false
 
 # Process and display mathematical expressions. This can clash with "$" characters in messages.
-latex = false
+latex = true
 
 # Authorize users to upload files with messages
 multi_modal = true

code/config.yml

@@ -2,14 +2,14 @@ embedding_options:
   embedd_files: False # bool
   data_path: 'storage/data' # str
   url_file_path: 'storage/data/urls.txt' # str
-  expand_urls: True # bool
-  db_option : 'FAISS' # str
+  expand_urls: False # bool
+  db_option : 'RAGatouille' # str [FAISS, Chroma, RAGatouille]
   db_path : 'vectorstores' # str
   model : 'sentence-transformers/all-MiniLM-L6-v2' # str [sentence-transformers/all-MiniLM-L6-v2, text-embedding-ada-002']
   search_top_k : 3 # int
-  score_threshold : 0.5 # float
+  score_threshold : 0.2 # float
 llm_params: 
-  use_history: True # bool
+  use_history: False # bool
   memory_window: 3 # int
   llm_loader: 'local_llm' # str [local_llm, openai]
   openai_params:

code/modules/data_loader.py

@@ -2,7 +2,7 @@ import os
 import re
 import requests
 import pysrt
-from langchain.document_loaders import (
+from langchain_community.document_loaders import (
     PyMuPDFLoader,
     Docx2txtLoader,
     YoutubeLoader,
@@ -16,6 +16,15 @@ import logging
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_experimental.text_splitter import SemanticChunker
 from langchain_openai.embeddings import OpenAIEmbeddings
+from ragatouille import RAGPretrainedModel
+from langchain.chains import LLMChain
+from langchain.llms import OpenAI
+from langchain import PromptTemplate
+
+try:
+    from modules.helpers import get_lecture_metadata
+except:
+    from helpers import get_lecture_metadata
 
 logger = logging.getLogger(__name__)
 
@@ -58,23 +67,6 @@ class FileReader:
             return None
 
     def read_pdf(self, temp_file_path: str):
-        # parser = LlamaParse(
-        #     api_key="",
-        #     result_type="markdown",
-        #     num_workers=4,
-        #     verbose=True,
-        #     language="en",
-        # )
-        # documents = parser.load_data(temp_file_path)
-
-        # with open("temp/output.md", "a") as f:
-        #     for doc in documents:
-        #         f.write(doc.text + "\n")
-
-        # markdown_path = "temp/output.md"
-        # loader = UnstructuredMarkdownLoader(markdown_path)
-        # loader = PyMuPDFLoader(temp_file_path)  # This loader preserves more metadata
-        # return loader.load()
         loader = self.pdf_reader.get_loader(temp_file_path)
         documents = self.pdf_reader.get_documents(loader)
         return documents
@@ -108,8 +100,6 @@ class FileReader:
 class ChunkProcessor:
     def __init__(self, config):
         self.config = config
-        self.document_chunks_full = []
-        self.document_names = []
 
         if config["splitter_options"]["use_splitter"]:
             if config["splitter_options"]["split_by_token"]:
@@ -130,6 +120,17 @@ class ChunkProcessor:
             self.splitter = None
         logger.info("ChunkProcessor instance created")
 
+    # def extract_metadata(self, document_content):
+
+    #     llm = OpenAI()
+    #     prompt_template = PromptTemplate(
+    #         input_variables=["document_content"],
+    #         template="Extract metadata for this document:\n\n{document_content}\n\nMetadata:",
+    #     )
+    #     chain = LLMChain(llm=llm, prompt=prompt_template)
+    #     metadata = chain.run(document_content=document_content)
+    #     return metadata
+
     def remove_delimiters(self, document_chunks: list):
         for chunk in document_chunks:
             for delimiter in self.config["splitter_options"]["delimiters_to_remove"]:
@@ -146,11 +147,23 @@ class ChunkProcessor:
         logger.info(f"\tNumber of pages after skipping: {len(document_chunks)}")
         return document_chunks
 
-    def process_chunks(self, documents):
-        if self.splitter:
+    def process_chunks(
+        self, documents, file_type="txt", source="", page=0, metadata={}
+    ):
+        documents = [Document(page_content=documents, source=source, page=page)]
+        if file_type == "txt":
             document_chunks = self.splitter.split_documents(documents)
-        else:
-            document_chunks = documents
+        elif file_type == "pdf":
+            document_chunks = documents  # Full page for now
+
+        # add the source and page number back to the metadata
+        for chunk in document_chunks:
+            chunk.metadata["source"] = source
+            chunk.metadata["page"] = page
+
+            # add the metadata extracted from the document
+            for key, value in metadata.items():
+                chunk.metadata[key] = value
 
         if self.config["splitter_options"]["remove_leftover_delimiters"]:
             document_chunks = self.remove_delimiters(document_chunks)
@@ -161,38 +174,77 @@ class ChunkProcessor:
 
     def get_chunks(self, file_reader, uploaded_files, weblinks):
         self.document_chunks_full = []
-        self.document_names = []
+        self.parent_document_names = []
+        self.child_document_names = []
+        self.documents = []
+        self.document_metadata = []
+
+        lecture_metadata = get_lecture_metadata(
+            "https://dl4ds.github.io/sp2024/lectures/"
+        )  # TODO: Use more efficiently
 
         for file_index, file_path in enumerate(uploaded_files):
             file_name = os.path.basename(file_path)
             file_type = file_name.split(".")[-1].lower()
 
-            try:
-                if file_type == "pdf":
-                    documents = file_reader.read_pdf(file_path)
-                elif file_type == "txt":
-                    documents = file_reader.read_txt(file_path)
-                elif file_type == "docx":
-                    documents = file_reader.read_docx(file_path)
-                elif file_type == "srt":
-                    documents = file_reader.read_srt(file_path)
+            # try:
+            if file_type == "pdf":
+                documents = file_reader.read_pdf(file_path)
+            elif file_type == "txt":
+                documents = file_reader.read_txt(file_path)
+            elif file_type == "docx":
+                documents = file_reader.read_docx(file_path)
+            elif file_type == "srt":
+                documents = file_reader.read_srt(file_path)
+            else:
+                logger.warning(f"Unsupported file type: {file_type}")
+                continue
+
+            # full_text = ""
+            # for doc in documents:
+            #     full_text += doc.page_content
+            #     break  # getting only first page for now
+
+            # extracted_metadata = self.extract_metadata(full_text)
+
+            for doc in documents:
+                page_num = doc.metadata.get("page", 0)
+                self.documents.append(doc.page_content)
+                self.document_metadata.append({"source": file_path, "page": page_num})
+                if "lecture" in file_path.lower():
+                    metadata = lecture_metadata.get(file_path, {})
+                    metadata["source_type"] = "lecture"
+                    self.document_metadata[-1].update(metadata)
                 else:
-                    logger.warning(f"Unsupported file type: {file_type}")
-                    continue
-
-                document_chunks = self.process_chunks(documents)
-                self.document_names.append(file_name)
-                self.document_chunks_full.extend(document_chunks)
+                    metadata = {"source_type": "other"}
+
+                self.child_document_names.append(f"{file_name}_{page_num}")
+
+                self.parent_document_names.append(file_name)
+                if self.config["embedding_options"]["db_option"] not in ["RAGatouille"]:
+                    document_chunks = self.process_chunks(
+                        self.documents[-1],
+                        file_type,
+                        source=file_path,
+                        page=page_num,
+                        metadata=metadata,
+                    )
+                    self.document_chunks_full.extend(document_chunks)
 
-            except Exception as e:
-                logger.error(f"Error processing file {file_name}: {str(e)}")
+            # except Exception as e:
+            #     logger.error(f"Error processing file {file_name}: {str(e)}")
 
         self.process_weblinks(file_reader, weblinks)
 
         logger.info(
             f"Total document chunks extracted: {len(self.document_chunks_full)}"
         )
-        return self.document_chunks_full, self.document_names
+        return (
+            self.document_chunks_full,
+            self.child_document_names,
+            self.documents,
+            self.document_metadata,
+        )
 
     def process_weblinks(self, file_reader, weblinks):
         if weblinks[0] != "":
@@ -206,9 +258,26 @@ class ChunkProcessor:
                     else:
                         documents = file_reader.read_html(link)
 
-                    document_chunks = self.process_chunks(documents)
-                    self.document_names.append(link)
-                    self.document_chunks_full.extend(document_chunks)
+                    for doc in documents:
+                        page_num = doc.metadata.get("page", 0)
+                        self.documents.append(doc.page_content)
+                        self.document_metadata.append(
+                            {"source": link, "page": page_num}
+                        )
+                        self.child_document_names.append(f"{link}")
+
+                    self.parent_document_names.append(link)
+                    if self.config["embedding_options"]["db_option"] not in [
+                        "RAGatouille"
+                    ]:
+                        document_chunks = self.process_chunks(
+                            self.documents[-1],
+                            "txt",
+                            source=link,
+                            page=0,
+                            metadata={"source_type": "webpage"},
+                        )
+                        self.document_chunks_full.extend(document_chunks)
                 except Exception as e:
                     logger.error(
                         f"Error splitting link {link_index+1} : {link}: {str(e)}"

code/modules/embedding_model_loader.py

@@ -1,6 +1,6 @@
 from langchain_community.embeddings import OpenAIEmbeddings
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.embeddings import LlamaCppEmbeddings
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.embeddings import LlamaCppEmbeddings
 
 try:
     from modules.constants import *

code/modules/helpers.py

@@ -4,6 +4,8 @@ from tqdm import tqdm
 from urllib.parse import urlparse
 import chainlit as cl
 from langchain import PromptTemplate
+import requests
+from bs4 import BeautifulSoup
 
 try:
     from modules.constants import *
@@ -138,67 +140,133 @@ def get_prompt(config):
 
 
 def get_sources(res, answer):
-    source_elements_dict = {}
     source_elements = []
-    found_sources = []
-
     source_dict = {}  # Dictionary to store URL elements
 
     for idx, source in enumerate(res["source_documents"]):
         source_metadata = source.metadata
         url = source_metadata["source"]
         score = source_metadata.get("score", "N/A")
+        page = source_metadata.get("page", 1)
+
+        lecture_tldr = source_metadata.get("tldr", "N/A")
+        lecture_recording = source_metadata.get("lecture_recording", "N/A")
+        suggested_readings = source_metadata.get("suggested_readings", "N/A")
 
-        if url not in source_dict:
-            source_dict[url] = [(source.page_content, score)]
+        source_type = source_metadata.get("source_type", "N/A")
+
+        url_name = f"{url}_{page}"
+        if url_name not in source_dict:
+            source_dict[url_name] = {
+                "text": source.page_content,
+                "url": url,
+                "score": score,
+                "page": page,
+                "lecture_tldr": lecture_tldr,
+                "lecture_recording": lecture_recording,
+                "suggested_readings": suggested_readings,
+                "source_type": source_type,
+            }
         else:
-            source_dict[url].append((source.page_content, score))
+            source_dict[url_name]["text"] += f"\n\n{source.page_content}"
 
-    for source_idx, (url, text_list) in enumerate(source_dict.items()):
-        full_text = ""
-        for url_idx, (text, score) in enumerate(text_list):
-            full_text += f"Source {url_idx + 1} (Score: {score}):\n{text}\n\n\n"
-        source_elements.append(cl.Text(name=url, content=full_text))
-        found_sources.append(f"{url} (Score: {score})")
+    # First, display the answer
+    full_answer = "**Answer:**\n"
+    full_answer += answer
 
-    if found_sources:
-        answer += f"\n\nSources: {', '.join(found_sources)}"
-    else:
-        answer += f"\n\nNo source found."
-
-    # for idx, source in enumerate(res["source_documents"]):
-    #     title = source.metadata["source"]
-
-    #     if title not in source_elements_dict:
-    #         source_elements_dict[title] = {
-    #             "page_number": [source.metadata["page"]],
-    #             "url": source.metadata["source"],
-    #             "content": source.page_content,
-    #         }
-
-    #     else:
-    #         source_elements_dict[title]["page_number"].append(source.metadata["page"])
-    #     source_elements_dict[title][
-    #         "content_" + str(source.metadata["page"])
-    #     ] = source.page_content
-    #     # sort the page numbers
-    #     # source_elements_dict[title]["page_number"].sort()
-
-    # for title, source in source_elements_dict.items():
-    #     # create a string for the page numbers
-    #     page_numbers = ", ".join([str(x) for x in source["page_number"]])
-    #     text_for_source = f"Page Number(s): {page_numbers}\nURL: {source['url']}"
-    #     source_elements.append(cl.Pdf(name="File", path=title))
-    #     found_sources.append("File")
-    #     # for pn in source["page_number"]:
-    #     #     source_elements.append(
-    #     #         cl.Text(name=str(pn), content=source["content_"+str(pn)])
-    #     #     )
-    #     #     found_sources.append(str(pn))
-
-    # if found_sources:
-    #     answer += f"\nSource:{', '.join(found_sources)}"
-    # else:
-    #     answer += f"\nNo source found."
-
-    return answer, source_elements
+    # Then, display the sources
+    full_answer += "\n\n**Sources:**\n"
+    for idx, (url_name, source_data) in enumerate(source_dict.items()):
+        full_answer += f"\nSource {idx + 1} (Score: {source_data['score']}): {source_data['url']}\n"
+
+        name = f"Source {idx + 1} Text\n"
+        full_answer += name
+        source_elements.append(cl.Text(name=name, content=source_data["text"]))
+
+        # Add a PDF element if the source is a PDF file
+        if source_data["url"].lower().endswith(".pdf"):
+            name = f"Source {idx + 1} PDF\n"
+            full_answer += name
+            pdf_url = f"{source_data['url']}#page={source_data['page']+1}"
+            source_elements.append(cl.Pdf(name=name, url=pdf_url))
+
+    # Finally, include lecture metadata for each unique source
+    # displayed_urls = set()
+    # full_answer += "\n**Metadata:**\n"
+    # for url_name, source_data in source_dict.items():
+    #     if source_data["url"] not in displayed_urls:
+    #         full_answer += f"\nSource: {source_data['url']}\n"
+    #         full_answer += f"Type: {source_data['source_type']}\n"
+    #         full_answer += f"TL;DR: {source_data['lecture_tldr']}\n"
+    #         full_answer += f"Lecture Recording: {source_data['lecture_recording']}\n"
+    #         full_answer += f"Suggested Readings: {source_data['suggested_readings']}\n"
+    #         displayed_urls.add(source_data["url"])
+    full_answer += "\n**Metadata:**\n"
+    for url_name, source_data in source_dict.items():
+        full_answer += f"\nSource: {source_data['url']}\n"
+        full_answer += f"Page: {source_data['page']}\n"
+        full_answer += f"Type: {source_data['source_type']}\n"
+        full_answer += f"TL;DR: {source_data['lecture_tldr']}\n"
+        full_answer += f"Lecture Recording: {source_data['lecture_recording']}\n"
+        full_answer += f"Suggested Readings: {source_data['suggested_readings']}\n"
+
+    return full_answer, source_elements
+
+
+def get_lecture_metadata(schedule_url):
+    """
+    Function to get the lecture metadata from the schedule URL.
+    """
+    lecture_metadata = {}
+
+    # Get the main schedule page content
+    r = requests.get(schedule_url)
+    soup = BeautifulSoup(r.text, "html.parser")
+
+    # Find all lecture blocks
+    lecture_blocks = soup.find_all("div", class_="lecture-container")
+
+    for block in lecture_blocks:
+        try:
+            # Extract the lecture title
+            title = block.find("span", style="font-weight: bold;").text.strip()
+
+            # Extract the TL;DR
+            tldr = block.find("strong", text="tl;dr:").next_sibling.strip()
+
+            # Extract the link to the slides
+            slides_link_tag = block.find("a", title="Download slides")
+            slides_link = slides_link_tag["href"].strip() if slides_link_tag else None
+
+            # Extract the link to the lecture recording
+            recording_link_tag = block.find("a", title="Download lecture recording")
+            recording_link = (
+                recording_link_tag["href"].strip() if recording_link_tag else None
+            )
+
+            # Extract suggested readings or summary if available
+            suggested_readings_tag = block.find("p", text="Suggested Readings:")
+            if suggested_readings_tag:
+                suggested_readings = suggested_readings_tag.find_next_sibling("ul")
+                if suggested_readings:
+                    suggested_readings = suggested_readings.get_text(
+                        separator="\n"
+                    ).strip()
+                else:
+                    suggested_readings = "No specific readings provided."
+            else:
+                suggested_readings = "No specific readings provided."
+
+            # Add to the dictionary
+            slides_link = f"https://dl4ds.github.io{slides_link}"
+            lecture_metadata[slides_link] = {
+                "tldr": tldr,
+                "title": title,
+                "lecture_recording": recording_link,
+                "suggested_readings": suggested_readings,
+            }
+        except Exception as e:
+            print(f"Error processing block: {e}")
+            continue
+
+    return lecture_metadata

code/modules/llm_tutor.py

@@ -8,7 +8,6 @@ from langchain.llms import CTransformers
 from langchain.memory import ConversationBufferWindowMemory
 from langchain.chains.conversational_retrieval.prompts import QA_PROMPT
 import os
-
 from modules.constants import *
 from modules.helpers import get_prompt
 from modules.chat_model_loader import ChatModelLoader
@@ -34,14 +33,21 @@ class LLMTutor:
 
     # Retrieval QA Chain
     def retrieval_qa_chain(self, llm, prompt, db):
-        retriever = VectorDBScore(
-            vectorstore=db,
-            search_type="similarity_score_threshold",
-            search_kwargs={
-                "score_threshold": self.config["embedding_options"]["score_threshold"],
-                "k": self.config["embedding_options"]["search_top_k"],
-            },
-        )
+        if self.config["embedding_options"]["db_option"] in ["FAISS", "Chroma"]:
+            retriever = VectorDBScore(
+                vectorstore=db,
+                search_type="similarity_score_threshold",
+                search_kwargs={
+                    "score_threshold": self.config["embedding_options"][
+                        "score_threshold"
+                    ],
+                    "k": self.config["embedding_options"]["search_top_k"],
+                },
+            )
+        elif self.config["embedding_options"]["db_option"] == "RAGatouille":
+            retriever = db.as_langchain_retriever(
+                k=self.config["embedding_options"]["search_top_k"]
+            )
         if self.config["llm_params"]["use_history"]:
             memory = ConversationBufferWindowMemory(
                 k=self.config["llm_params"]["memory_window"],

code/modules/vector_db.py

@@ -1,11 +1,12 @@
 import logging
 import os
 import yaml
-from langchain.vectorstores import FAISS, Chroma
+from langchain_community.vectorstores import FAISS, Chroma
 from langchain.schema.vectorstore import VectorStoreRetriever
 from langchain.callbacks.manager import CallbackManagerForRetrieverRun
 from langchain.schema.document import Document
 from langchain_core.callbacks import AsyncCallbackManagerForRetrieverRun
+from ragatouille import RAGPretrainedModel
 
 try:
     from modules.embedding_model_loader import EmbeddingModelLoader
@@ -25,7 +26,7 @@ class VectorDBScore(VectorStoreRetriever):
 
     # See https://github.com/langchain-ai/langchain/blob/61dd92f8215daef3d9cf1734b0d1f8c70c1571c3/libs/langchain/langchain/vectorstores/base.py#L500
     def _get_relevant_documents(
-            self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
     ) -> List[Document]:
         docs_and_similarities = (
             self.vectorstore.similarity_search_with_relevance_scores(
@@ -55,7 +56,6 @@ class VectorDBScore(VectorStoreRetriever):
         return docs
 
 
-
 class VectorDB:
     def __init__(self, config, logger=None):
         self.config = config
@@ -116,7 +116,15 @@ class VectorDB:
         self.embedding_model_loader = EmbeddingModelLoader(self.config)
         self.embedding_model = self.embedding_model_loader.load_embedding_model()
 
-    def initialize_database(self, document_chunks: list, document_names: list):
+    def initialize_database(
+        self,
+        document_chunks: list,
+        document_names: list,
+        documents: list,
+        document_metadata: list,
+    ):
+        if self.db_option in ["FAISS", "Chroma"]:
+            self.create_embedding_model()
         # Track token usage
         self.logger.info("Initializing vector_db")
         self.logger.info("\tUsing {} as db_option".format(self.db_option))
@@ -136,6 +144,14 @@ class VectorDB:
                     + self.config["embedding_options"]["model"],
                 ),
             )
+        elif self.db_option == "RAGatouille":
+            self.RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")
+            index_path = self.RAG.index(
+                index_name="new_idx",
+                collection=documents,
+                document_ids=document_names,
+                document_metadatas=document_metadata,
+            )
         self.logger.info("Completed initializing vector_db")
 
     def create_database(self):
@@ -146,11 +162,13 @@ class VectorDB:
         files += lecture_pdfs
         if "storage/data/urls.txt" in files:
             files.remove("storage/data/urls.txt")
-        document_chunks, document_names = data_loader.get_chunks(files, urls)
+        document_chunks, document_names, documents, document_metadata = (
+            data_loader.get_chunks(files, urls)
+        )
         self.logger.info("Completed loading data")
-
-        self.create_embedding_model()
-        self.initialize_database(document_chunks, document_names)
+        self.initialize_database(
+            document_chunks, document_names, documents, document_metadata
+        )
 
     def save_database(self):
         if self.db_option == "FAISS":
@@ -166,6 +184,9 @@ class VectorDB:
         elif self.db_option == "Chroma":
             # db is saved in the persist directory during initialization
             pass
+        elif self.db_option == "RAGatouille":
+            # index is saved during initialization
+            pass
         self.logger.info("Saved database")
 
     def load_database(self):
@@ -180,7 +201,7 @@ class VectorDB:
                     + self.config["embedding_options"]["model"],
                 ),
                 self.embedding_model,
-                # allow_dangerous_deserialization=True, <- unexpected keyword argument to load_local
+                allow_dangerous_deserialization=True,
             )
         elif self.db_option == "Chroma":
             self.vector_db = Chroma(
@@ -193,6 +214,10 @@ class VectorDB:
                 ),
                 embedding_function=self.embedding_model,
             )
+        elif self.db_option == "RAGatouille":
+            self.vector_db = RAGPretrainedModel.from_index(
+                ".ragatouille/colbert/indexes/new_idx"
+            )
         self.logger.info("Loaded database")
         return self.vector_db
 

requirements.txt

@@ -17,3 +17,4 @@ fake-useragent==1.4.0
 git+https://github.com/huggingface/accelerate.git
 llama-cpp-python
 PyPDF2==3.0.1
+ragatouille==0.0.8.post2

storage/data/urls.txt

@@ -1 +1,3 @@
 https://dl4ds.github.io/sp2024/
+https://dl4ds.github.io/sp2024/static_files/lectures/15_RAG_CoT.pdf
+https://dl4ds.github.io/sp2024/static_files/lectures/21_RL_RLHF_v2.pdf