new code

src/app.py

@@ -1,19 +1,14 @@
 from interface import create_demo
 from pdfchatbot import PDFChatBot
 
-# Create Gradio interface
 demo, chat_history, show_img, txt, submit_button, uploaded_pdf, slider1 = create_demo()
 
-# Create PDFChatBot instance
 pdf_chatbot = PDFChatBot()
 pdf_chatbot.create_organic_pipeline()
 
-# Set up event handlers
 with demo:
-    # Event handler for uploading a PDF
     uploaded_pdf.upload(pdf_chatbot.render_file, inputs=[uploaded_pdf,slider1], outputs=[show_img])
 
-    # Event handler for submitting text and generating response
     submit_button.click(pdf_chatbot.add_text, inputs=[chat_history, txt], outputs=[chat_history], queue=False).\
         success(pdf_chatbot.generate_response, inputs=[chat_history, txt, uploaded_pdf], outputs=[chat_history,txt]).\
         success(pdf_chatbot.render_file, inputs=[uploaded_pdf], outputs=[show_img])

src/pdfchatbot.py

@@ -40,62 +40,11 @@ class PDFChatBot:
         self.format_seperator="""\n\n--\n\n"""
         #self.chunk_size_slider = chunk_size_slider
 
-    def load_config(self, file_path):
-        """
-        Load configuration from a YAML file.
-
-        Parameters:
-            file_path (str): Path to the YAML configuration file.
-
-        Returns:
-            dict: Configuration as a dictionary.
-        """
-        with open(file_path, 'r') as stream:
-            try:
-                config = yaml.safe_load(stream)
-                return config
-            except yaml.YAMLError as exc:
-                print(f"Error loading configuration: {exc}")
-                return None
-
-    def add_text(self, history, text):
-        """
-        Add user-entered text to the chat history.
-
-        Parameters:
-            history (list): List of chat history tuples.
-            text (str): User-entered text.
-
-        Returns:
-            list: Updated chat history.
-        """
-        if not text:
-            raise gr.Error('Enter text')
-        history.append((text, ''))
-        #self.chunk_size_slider.interactive = False
-        return history
-
-    def create_prompt_template(self):
-        """
-        Create a prompt template for the chatbot.
-        """
-        template = (
-            f"The assistant should provide detailed explanations."
-            "Combine the chat history and follow up question into "
-            "Follow up question: What is this"
-        )
-        self.prompt = PromptTemplate.from_template(template)
-
     def load_embeddings(self):
-        """
-        Load embeddings from Hugging Face and set in the config file.
-        """
-        self.embeddings = HuggingFaceEmbeddings(model_name=self.config.get("modelEmbeddings"))
+        self.embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        print("Embedding model loaded")
 
     def load_vectordb(self):
-        """
-        Load the vector database from the documents and embeddings.
-        """
         text_splitter = RecursiveCharacterTextSplitter(
             chunk_size=256,
             chunk_overlap=100,
@@ -104,49 +53,23 @@ class PDFChatBot:
         )
         docs = text_splitter.split_documents(self.documents)
         self.vectordb = Chroma.from_documents(docs, self.embeddings)
-
-    def load_tokenizer(self):
-        """
-        Load the tokenizer from Hugging Face and set in the config file.
-        """
-        self.tokenizer = AutoTokenizer.from_pretrained(self.config.get("autoTokenizer"))
-
-    def load_model(self):
-        """
-        Load the causal language model from Hugging Face and set in the config file.
-        """
-        self.model = AutoModelForCausalLM.from_pretrained(
-            self.config.get("autoModelForCausalLM"),
-            device_map='auto',
-            torch_dtype=torch.float16,
-            token=True,
-        )
-
-    def create_pipeline(self):
-        """
-        Create a pipeline for text generation using the loaded model and tokenizer.
-        """
-        pipe = pipeline(
-            model=self.model,
-            task='text-generation',
-            tokenizer=self.tokenizer,
-            max_new_tokens=1024
-        )
-        self.pipeline = HuggingFacePipeline(pipeline=pipe)
+        print("Vector store created")
 
     def create_organic_pipeline(self):
         self.pipeline = pipeline(
             "text-generation",
             model="meta-llama/Meta-Llama-3-8B-Instruct",
             tokenizer="meta-llama/Meta-Llama-3-8B-Instruct",
-            model_kwargs={"torch_dtype": torch.bfloat16},
+            model_kwargs={"torch_dtype": torch.float16},
             device="cuda",
         )
+        print("Model pipeline loaded")
 
     def get_organic_context(self, query):
         documents = self.vectordb.similarity_search_with_relevance_scores(query, k=3)
         context = self.format_seperator.join([doc.page_content for doc, score in documents])
         self.current_context = context
+        print("Context Ready")
         print(self.current_context)
 
     def create_organic_response(self, history, query):
@@ -175,17 +98,6 @@ class PDFChatBot:
         )
         return outputs[0]["generated_text"][len(prompt):]
 
-    def create_chain(self):
-        """
-        Create a Conversational Retrieval Chain
-        """
-        self.chain = ConversationalRetrievalChain.from_llm(
-            self.pipeline,
-            chain_type="stuff",
-            retriever=self.vectordb.as_retriever(search_kwargs={"k": 1}),
-            condense_question_prompt=self.prompt,
-            return_source_documents=True
-        )
 
     def process_file(self, file):
         """
@@ -194,24 +106,13 @@ class PDFChatBot:
         Parameters:
             file (FileStorage): The uploaded PDF file.
         """
-        self.create_prompt_template()
         self.documents = PyPDFLoader(file.name).load()
         self.load_embeddings()
         self.load_vectordb()
         #self.create_chain()
     @spaces.GPU
     def generate_response(self, history, query, file):
-        """
-        Generate a response based on user query and chat history.
 
-        Parameters:
-            history (list): List of chat history tuples.
-            query (str): User's query.
-            file (FileStorage): The uploaded PDF file.
-
-        Returns:
-            tuple: Updated chat history and a space.
-        """
         if not query:
             raise gr.Error(message='Submit a question')
         if not file:
@@ -220,31 +121,28 @@ class PDFChatBot:
             self.process_file(file)
             self.processed = True
 
-        """result = self.chain({"question": query, 'chat_history': self.chat_history}, return_only_outputs=True)
-        self.chat_history.append((query, result["answer"]))
-        for char in result['answer']:
-            history[-1][-1] += char
-        return history, " """""
-
         result = self.create_organic_response(history="",query=query)
         return result,""
 
     def render_file(self, file,chunk_size):
-        """
-        Renders a specific page of a PDF file as an image.
-
-        Parameters:
-            file (FileStorage): The PDF file.
-
-        Returns:
-            PIL.Image.Image: The rendered page as an image.
-        """
-        #if self.chunk_size_slider.interactive == False:
-        #    self.chunk_size_slider.interactive = True
         print(chunk_size)
         doc = fitz.open(file.name)
         page = doc[self.page]
         self.chunk_size = chunk_size
         pix = page.get_pixmap(matrix=fitz.Matrix(300 / 72, 300 / 72))
         image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
-        return image
+        return image
+
+    def add_text(self, history, text):
+        """
+        Add user-entered text to the chat history.
+        Parameters:
+            history (list): List of chat history tuples.
+            text (str): User-entered text.
+        Returns:
+            list: Updated chat history.
+        """
+        if not text:
+            raise gr.Error('Enter text')
+        history.append((text, ''))
+        return history