Merge branch 'rag'

app.py

@@ -23,7 +23,9 @@ with gr.Blocks(fill_height=True) as demo:
         with gr.Column():
             url = gr.Textbox(value="https://www.gradio.app/docs/gradio/chatinterface", label="Docs URL", render=True)
             chat = gr.ChatInterface(
-                model.respond,
+                # model.respond,
+                model.predict,
+                # model.rag,
                 additional_inputs=[
                     url,
                     max_tokens,

data.py

@@ -1,26 +1,34 @@
 from langchain_community.document_loaders import UnstructuredURLLoader
 
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import PyPDFDirectoryLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
+
 class Data:
-    def __init__(self, url):
-        self.url = url
-
-    def get_context(self):
-        urls = [
-            self.url,
-        ]
-        loader = UnstructuredURLLoader(urls=urls)
+    def __init__(self, urls):
+        self.urls = urls
+         ## Embedding Using Huggingface
+        self.huggingface_embeddings = HuggingFaceBgeEmbeddings(
+            model_name="BAAI/bge-small-en-v1.5",      #sentence-transformers/all-MiniLM-l6-v2
+            model_kwargs={'device':'cpu'},
+            encode_kwargs={'normalize_embeddings':True}
+        )
+
+    @property
+    def retriever(self):
+        loader = UnstructuredURLLoader(urls=self.urls)
         data = loader.load()
-    
-        context = data[0].page_content # will come from 'url'
-
-        return context
-
-    def build_prompt(self, question):
-        prompt = f"""
-        Use the following piece of context to answer the question asked.
-        Please try to provide the answer only based on the context
-        {self.get_context()}
-        Question:{question}
-        Helpful Answers:
-        """
-        return prompt
+
+        ## VectorStore Creation
+        vectorstore = FAISS.from_documents(data, self.huggingface_embeddings)
+
+        retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k":3})
+
+        return retriever
+      
+
+        

model.py

@@ -1,12 +1,76 @@
 import os
 
 from huggingface_hub import InferenceClient
+from langchain.schema import SystemMessage, AIMessage, HumanMessage
+
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 
 from data import Data
 
 class Model:
     def __init__(self, model_id="meta-llama/Llama-3.2-1B-Instruct"):
         self.client = InferenceClient(model_id, token=os.getenv("HUGGINGFACEHUB_API_TOKEN"))
+        self.llm = HuggingFaceEndpoint(
+            repo_id="HuggingFaceH4/zephyr-7b-beta",
+            task="text-generation",
+            max_new_tokens=512,
+            do_sample=False,
+            repetition_penalty=1.03,
+        )
+        self.chat_model = ChatHuggingFace(llm=self.llm, token=os.getenv("HUGGINGFACEHUB_API_TOKEN"))
+
+    def build_prompt(self, question, context_urls):
+        data = Data(context_urls)
+        context = data.retriever.invoke(f"{question}")[0].page_content
+        prompt = f"""
+        Use the following piece of context to answer the question asked.
+        Please try to provide the answer only based on the context
+        {context}
+        Question:{question}
+        Helpful Answers:
+        """
+        return prompt
+
+    def _build_prompt_rag(self):
+        prompt_template="""
+        Use the following piece of context to answer the question asked.
+        Please try to provide the answer only based on the context
+        {context}
+        Question:{question}
+        Helpful Answers:
+         """
+        prompt=PromptTemplate(template=prompt_template,input_variables=["context","question"])
+        return prompt
+
+    def _retrieval_qa(self, url):
+        data = Data([url])
+        prompt = self._build_prompt_rag()
+        return RetrievalQA.from_chain_type(
+            llm=self.chat_model,
+            chain_type="stuff",
+            retriever=data.retriever,
+            return_source_documents=True,
+            chain_type_kwargs={"prompt":prompt}
+        )
+
+    def predict(self, message, history, url, max_tokens, temperature, top_p):
+        history_langchain_format = [SystemMessage(content="You're a helpful python developer assistant")]
+        for msg in history:
+            if msg['role'] == "user":
+                history_langchain_format.append(HumanMessage(content=msg['content']))
+            elif msg['role'] == "assistant":
+                history_langchain_format.append(AIMessage(content=msg['content']))
+        history_langchain_format.append(HumanMessage(content=message))
+
+        # ai_msg = self.chat_model.invoke(history_langchain_format)
+        # return ai_msg.content
+
+        ret = self._retrieval_qa(url)
+        return ret.invoke({"query": message})['result']
+
 
     def respond(
         self,
@@ -17,9 +81,7 @@ class Model:
         temperature,
         top_p,
     ):
-    
-        data = Data(url)
-    
+        
         messages = [{"role": "system", "content": url}]
     
         for val in history:
@@ -28,7 +90,7 @@ class Model:
             if val[1]:
                 messages.append({"role": "assistant", "content": val[1]})
     
-        messages.append({"role": "user", "content": data.build_prompt(message)})
+        messages.append({"role": "user", "content": self.build_prompt(message, [url])})
     
         response = ""
     

requirements.txt

@@ -1,4 +1,7 @@
 huggingface_hub==0.25.2
 langchain-community==0.3.3
+langchain-core==0.3.12
+langchain-huggingface==0.1.0
 unstructured==0.16.0
-unstructured-client==0.26.1
+unstructured-client==0.26.1
+faiss-cpu==1.9.0