Formatting

app.py

@@ -15,23 +15,24 @@ llm_model_name = "gpt-4o-mini"
 # Settting up LLMGenerator
 llm_generator = None
 
+
 def set_api_key(api_key: str):
     if api_key.strip():
-        os.environ['OPENAI_API_KEY'] = api_key
+        os.environ["OPENAI_API_KEY"] = api_key
     else:
         raise gr.Error("Please provide a valid API key")
 
-def process_inputs(api_key:str, pdf_file, questions: str):
+
+def process_inputs(api_key: str, pdf_file, questions: str):
 
     # Setup Api KEY
     set_api_key(api_key)
-    
 
     if pdf_file is None:
         raise gr.Error("Please upload a pdf file")
-    
+
     # Parsing the pdf
-    doc_handler = DocParsing(file_path=pdf_file.name,model_name=embedding_model_name)
+    doc_handler = DocParsing(file_path=pdf_file.name, model_name=embedding_model_name)
     docs = doc_handler.process_pdf()
 
     # Create vector store
@@ -43,26 +44,32 @@ def process_inputs(api_key:str, pdf_file, questions: str):
     if not questions.strip():
         raise gr.Error("Please provide valid set of questions")
     output_dict = {}
-    questions_list = questions.strip().split('\n')
+    questions_list = questions.strip().split("\n")
     for question in questions_list:
-        
+
         # Retrieve top similar chunks
         similar_chunks = retriever.search(query=question, k=10)
 
         # Generate the answer
         output_dict[question] = llm_generator.generate_answer(question, similar_chunks)
-        
-    
+
     response = json.dumps(output_dict, indent=4)
     return response
 
+
 with gr.Blocks() as demo:
     gr.Markdown("# AskMYPDF Q&A App")
-    gr.Markdown("Enter your OPENAI API key, upload a PDF, and list your questions below.")
-    
+    gr.Markdown(
+        "Enter your OPENAI API key, upload a PDF, and list your questions below."
+    )
+
     api_key_input = gr.Textbox(label="API Key", type="password")
     pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
-    questions_input = gr.Textbox(label="List of Questions (one per line)", lines=5, placeholder="Question 1\nQuestion 2\n...")
+    questions_input = gr.Textbox(
+        label="List of Questions (one per line)",
+        lines=5,
+        placeholder="Question 1\nQuestion 2\n...",
+    )
 
     submit_button = gr.Button("Submit")
     output = gr.Textbox(label="Output")
@@ -70,8 +77,8 @@ with gr.Blocks() as demo:
     submit_button.click(
         fn=process_inputs,
         inputs=[api_key_input, pdf_input, questions_input],
-        outputs=output
+        outputs=output,
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

utils/llm_generation.py

@@ -1,6 +1,5 @@
 from langchain.prompts import (
     ChatPromptTemplate,
-    SystemMessagePromptTemplate,
     HumanMessagePromptTemplate,
 )
 from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
@@ -57,9 +56,7 @@ class LLMGeneration:
                 content="<context>John traveled to Paris last summer. He stayed at a small boutique hotel and visited the Louvre museum.</context>"
                 "<question>Where did John travel?</question>"
             ),
-            AIMessage(
-                content="""{"answer": "Paris"}"""
-            ),
+            AIMessage(content="""{"answer": "Paris"}"""),
         ]
 
         self.initial_prompt_messages = [system_message] + few_shots

utils/retrieval.py

@@ -3,11 +3,14 @@ from langchain_community.vectorstores import FAISS
 from langchain.schema import Document
 from typing import List
 
+
 class Retrieval:
-    def __init__(self, model_name):
+    def __init__(self, model_name, max_model_tokens=384):
         self.model_name = model_name
-        self.embeddings = HuggingFaceEmbeddings(model_name=model_name)
-        
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name=model_name,
+            encode_kwargs={"max_length": max_model_tokens, "truncation": True},
+        )
 
     def create_vector_store(self, chunks: List[Document]):
 
@@ -15,8 +18,8 @@ class Retrieval:
         # Create FAISS vector store
         self.vectorstore = FAISS.from_documents(self.chunks, self.embeddings)
 
-    def search(self,query, k=10) -> List[Document]:
+    def search(self, query, k=10) -> List[Document]:
         # Retrieve top 10 similar chunks
         similar_docs = self.vectorstore.similarity_search(query, k)
 
-        return similar_docs
+        return similar_docs