Upload 5 files

black.py

@@ -0,0 +1,131 @@
+import os
+
+os.environ["HF_HOME"] = "weights"
+os.environ["TORCH_HOME"] = "weights"
+
+import gc
+import re
+import uuid
+import textwrap
+import subprocess
+import nest_asyncio
+from dotenv import load_dotenv
+from IPython.display import Markdown, display
+
+from llama_index.core import Settings
+from llama_index.llms.ollama import Ollama
+from llama_index.core import PromptTemplate
+from llama_index.core import SimpleDirectoryReader
+from llama_index.core.ingestion import IngestionPipeline
+from llama_index.core import VectorStoreIndex
+from llama_index.core.storage.storage_context import StorageContext
+
+from langchain.embeddings import HuggingFaceEmbeddings
+from llama_index.embeddings.langchain import LangchainEmbedding
+
+from rag_101.retriever import (
+    load_embedding_model,
+    load_reranker_model
+)
+
+# allows nested access to the event loop
+nest_asyncio.apply()
+
+# setting up the llm
+llm=Ollama(model="mistral", request_timeout=60.0)
+
+# setting up the embedding model
+lc_embedding_model = load_embedding_model()
+embed_model = LangchainEmbedding(lc_embedding_model)
+
+# utility functions
+def parse_github_url(url):
+    pattern = r"https://github\.com/([^/]+)/([^/]+)"
+    match = re.match(pattern, url)
+    return match.groups() if match else (None, None)
+
+def clone_github_repo(repo_url):    
+    try:
+        print('Cloning the repo ...')
+        result = subprocess.run(["git", "clone", repo_url], check=True, text=True, capture_output=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to clone repository: {e}")
+        return None
+
+
+def validate_owner_repo(owner, repo):
+    return bool(owner) and bool(repo)
+
+# Setup a query engine
+
+def setup_query_engine(github_url):
+    
+    owner, repo = parse_github_url(github_url)
+    
+    if validate_owner_repo(owner, repo):
+        # Clone the GitHub repo & save it in a directory
+        input_dir_path = f"{repo}"
+
+        if os.path.exists(input_dir_path):
+            pass
+        else:
+            clone_github_repo(github_url)
+        
+        loader = SimpleDirectoryReader(
+            input_dir = input_dir_path,
+            required_exts=[".py", ".ipynb", ".js", ".ts", ".md"],
+            recursive=True
+        )
+
+        try:
+            docs = loader.load_data()
+
+            # ====== Create vector store and upload data ======
+            Settings.embed_model = embed_model
+            index = VectorStoreIndex.from_documents(docs, show_progress=True)
+            # TODO try async index creation for faster emebdding generation & persist it to memory!
+            # index = VectorStoreIndex(docs, use_async=True)
+
+            # ====== Setup a query engine ======
+            Settings.llm = llm
+            query_engine = index.as_query_engine(similarity_top_k=4)
+            
+            # ====== Customise prompt template ======
+            qa_prompt_tmpl_str = (
+            "Context information is below.\n"
+            "---------------------\n"
+            "{context_str}\n"
+            "---------------------\n"
+            "Given the context information above I want you to think step by step to answer the query in a crisp manner, incase case you don't know the answer say 'I don't know!'.\n"
+            "Query: {query_str}\n"
+            "Answer: "
+            )
+            qa_prompt_tmpl = PromptTemplate(qa_prompt_tmpl_str)
+
+            query_engine.update_prompts(
+                {"response_synthesizer:text_qa_template": qa_prompt_tmpl}
+            )
+
+            if docs:
+                print("Data loaded successfully!!")
+                print("Ready to chat!!")
+            else:
+                print("No data found, check if the repository is not empty!")
+            
+            return query_engine
+
+        except Exception as e:
+            print(f"An error occurred: {e}")
+    else:
+        print('Invalid github repo, try again!')
+        return None
+
+# Provide url to the repository you want to chat with
+github_url = "https://github.com/Aniket23160/Pose-Graph-SLAM"
+
+query_engine = setup_query_engine(github_url=github_url)
+print("----------------------------------------------------------------")
+query='What is this repo about?'
+print(f"Question: {query}")
+response = query_engine.query(query)
+print(f"Answer: {response}")

rag_101/client.py

@@ -0,0 +1,52 @@
+import time
+from typing import List, Optional, Union
+
+from langchain_community.chat_models import ChatOllama
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from retriever import (
+    create_parent_retriever,
+    load_embedding_model,
+    load_pdf,
+    load_reranker_model,
+    retrieve_context,
+)
+
+
+def main(
+    file: str = "example_data/2401.08406.pdf",
+    llm_name="mistral",
+):
+    docs = load_pdf(files=file)
+
+    embedding_model = load_embedding_model()
+    retriever = create_parent_retriever(docs, embedding_model)
+    reranker_model = load_reranker_model()
+
+    llm = ChatOllama(model=llm_name)
+    prompt_template = ChatPromptTemplate.from_template(
+        (
+            "Please answer the following question based on the provided `context` that follows the question.\n"
+            "If you do not know the answer then just say 'I do not know'\n"
+            "question: {question}\n"
+            "context: ```{context}```\n"
+        )
+    )
+    chain = prompt_template | llm | StrOutputParser()
+
+    while True:
+        query = input("Ask question: ")
+        context = retrieve_context(
+            query, retriever=retriever, reranker_model=reranker_model
+        )[0]
+        print("LLM Response: ", end="")
+        for e in chain.stream({"context": context[0].page_content, "question": query}):
+            print(e, end="")
+        print()
+        time.sleep(0.1)
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    CLI(main)

rag_101/rag.py

@@ -0,0 +1,61 @@
+from langchain.callbacks import FileCallbackHandler
+from langchain_community.chat_models import ChatOllama
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from loguru import logger
+
+from rag_101.retriever import (
+    RAGException,
+    create_parent_retriever,
+    load_embedding_model,
+    load_pdf,
+    load_reranker_model,
+    retrieve_context,
+)
+
+
+class RAGClient:
+    embedding_model = load_embedding_model()
+    reranker_model = load_reranker_model()
+
+    def __init__(self, files, model="mistral"):
+        docs = load_pdf(files=files)
+        self.retriever = create_parent_retriever(docs, self.embedding_model)
+
+        llm = ChatOllama(model=model)
+        prompt_template = ChatPromptTemplate.from_template(
+            (
+                "Please answer the following question based on the provided `context` that follows the question.\n"
+                "Think step by step before coming to answer. If you do not know the answer then just say 'I do not know'\n"
+                "question: {question}\n"
+                "context: ```{context}```\n"
+            )
+        )
+        self.chain = prompt_template | llm | StrOutputParser()
+
+    def stream(self, query: str) -> dict:
+        try:
+            context, similarity_score = self.retrieve_context(query)[0]
+            context = context.page_content
+            if similarity_score < 0.005:
+                context = "This context is not confident. " + context
+        except RAGException as e:
+            context, similarity_score = e.args[0], 0
+        logger.info(context)
+        for r in self.chain.stream({"context": context, "question": query}):
+            yield r
+
+    def retrieve_context(self, query: str):
+        return retrieve_context(
+            query, retriever=self.retriever, reranker_model=self.reranker_model
+        )
+
+    def generate(self, query: str) -> dict:
+        contexts = self.retrieve_context(query)
+
+        return {
+            "contexts": contexts,
+            "response": self.chain.invoke(
+                {"context": contexts[0][0].page_content, "question": query}
+            ),
+        }

rag_101/retriever.py

@@ -0,0 +1,160 @@
+import os
+
+os.environ["HF_HOME"] = "weights"
+os.environ["TORCH_HOME"] = "weights"
+
+from typing import List, Optional, Union
+
+from langchain.callbacks import FileCallbackHandler
+from langchain.retrievers import ContextualCompressionRetriever, ParentDocumentRetriever
+from langchain.retrievers.document_compressors import EmbeddingsFilter
+from langchain.storage import InMemoryStore
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import UnstructuredFileLoader
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from langchain_community.vectorstores import FAISS, Chroma
+from langchain_core.documents import Document
+from loguru import logger
+from rich import print
+from sentence_transformers import CrossEncoder
+from unstructured.cleaners.core import clean_extra_whitespace, group_broken_paragraphs
+
+logfile = "log/output.log"
+logger.add(logfile, colorize=True, enqueue=True)
+handler = FileCallbackHandler(logfile)
+
+
+persist_directory = None
+
+
+class RAGException(Exception):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+
+def rerank_docs(reranker_model, query, retrieved_docs):
+    query_and_docs = [(query, r.page_content) for r in retrieved_docs]
+    scores = reranker_model.predict(query_and_docs)
+    return sorted(list(zip(retrieved_docs, scores)), key=lambda x: x[1], reverse=True)
+
+
+def load_pdf(
+    files: Union[str, List[str]] = "example_data/2401.08406.pdf"
+) -> List[Document]:
+    if isinstance(files, str):
+        loader = UnstructuredFileLoader(
+            files,
+            post_processors=[clean_extra_whitespace, group_broken_paragraphs],
+        )
+        return loader.load()
+
+    loaders = [
+        UnstructuredFileLoader(
+            file,
+            post_processors=[clean_extra_whitespace, group_broken_paragraphs],
+        )
+        for file in files
+    ]
+    docs = []
+    for loader in loaders:
+        docs.extend(
+            loader.load(),
+        )
+    return docs
+
+
+def create_parent_retriever(
+    docs: List[Document], embeddings_model: HuggingFaceBgeEmbeddings()
+):
+    parent_splitter = RecursiveCharacterTextSplitter(
+        separators=["\n\n\n", "\n\n"],
+        chunk_size=2000,
+        length_function=len,
+        is_separator_regex=False,
+    )
+
+    # This text splitter is used to create the child documents
+    child_splitter = RecursiveCharacterTextSplitter(
+        separators=["\n\n\n", "\n\n"],
+        chunk_size=1000,
+        chunk_overlap=300,
+        length_function=len,
+        is_separator_regex=False,
+    )
+    # The vectorstore to use to index the child chunks
+    vectorstore = Chroma(
+        collection_name="split_documents",
+        embedding_function=embeddings_model,
+        persist_directory=persist_directory,
+    )
+    # The storage layer for the parent documents
+    store = InMemoryStore()
+    retriever = ParentDocumentRetriever(
+        vectorstore=vectorstore,
+        docstore=store,
+        child_splitter=child_splitter,
+        parent_splitter=parent_splitter,
+        k=10,
+    )
+    retriever.add_documents(docs)
+    return retriever
+
+
+def retrieve_context(query, retriever, reranker_model):
+    retrieved_docs = retriever.get_relevant_documents(query)
+
+    if len(retrieved_docs) == 0:
+        raise RAGException(
+            f"Couldn't retrieve any relevant document with the query `{query}`. Try modifying your question!"
+        )
+    reranked_docs = rerank_docs(
+        query=query, retrieved_docs=retrieved_docs, reranker_model=reranker_model
+    )
+    return reranked_docs
+
+
+def load_embedding_model(
+    model_name: str = "BAAI/bge-large-en-v1.5", device: str = "cpu"
+) -> HuggingFaceBgeEmbeddings:
+    model_kwargs = {"device": device}
+    encode_kwargs = {
+        "normalize_embeddings": True
+    }  # set True to compute cosine similarity
+    embedding_model = HuggingFaceBgeEmbeddings(
+        model_name=model_name,
+        model_kwargs=model_kwargs,
+        encode_kwargs=encode_kwargs,
+    )
+    return embedding_model
+
+
+def load_reranker_model(
+    reranker_model_name: str = "BAAI/bge-reranker-large", device: str = "cpu"
+) -> CrossEncoder:
+    reranker_model = CrossEncoder(
+        model_name=reranker_model_name, max_length=512, device=device
+    )
+    return reranker_model
+
+
+def main(
+    file: str = "example_data/2401.08406.pdf",
+    query: Optional[str] = None,
+    llm_name="mistral",
+):
+    docs = load_pdf(files=file)
+
+    embedding_model = load_embedding_model()
+    retriever = create_parent_retriever(docs, embedding_model)
+    reranker_model = load_reranker_model()
+
+    context = retrieve_context(
+        query, retriever=retriever, reranker_model=reranker_model
+    )[0]
+    print("context:\n", context, "\n", "=" * 50, "\n")
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    CLI(main)