Update src/pdfchatbot.py

src/pdfchatbot.py

@@ -3,10 +3,10 @@ import fitz
 import torch
 import gradio as gr
 from PIL import Image
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
-from langchain_community.document_loaders import PyPDFLoader
+from langchain.document_loaders import PyPDFLoader
 from langchain.prompts import PromptTemplate
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import spaces
@@ -17,7 +17,6 @@ class PDFChatBot:
     def __init__(self, config_path="config.yaml"):
         """
         Initialize the PDFChatBot instance.
-
         Parameters:
             config_path (str): Path to the configuration file (default is "../config.yaml").
         """
@@ -33,9 +32,9 @@ class PDFChatBot:
         self.model = None
         self.pipeline = None
         self.chain = None
-        self.chunk_size = 16384
+        self.chunk_size = 512
         self.overlap_percentage = 50
-        self.max_chunks_in_context = 5
+        self.max_chunks_in_context = 2
         self.current_context = None
         self.model_temperatue = 0.5
         self.format_seperator="""\n\n--\n\n"""
@@ -60,13 +59,13 @@ class PDFChatBot:
         print("Vector store created")
     @spaces.GPU
     def load_tokenizer(self):
-        self.tokenizer = AutoTokenizer.from_pretrained("gradientai/Llama-3-8B-Instruct-Gradient-1048k")
+        self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
 
     @spaces.GPU
     def create_organic_pipeline(self):
         self.pipe = pipeline(
             "text-generation",
-            model="gradientai/Llama-3-8B-Instruct-Gradient-1048k",
+            model="meta-llama/Meta-Llama-3-8B-Instruct",
             model_kwargs={"torch_dtype": torch.bfloat16},
             device="cuda",
         )
@@ -84,7 +83,7 @@ class PDFChatBot:
         """
         pipe = pipeline(
             "text-generation",
-            model="gradientai/Llama-3-8B-Instruct-Gradient-1048k",
+            model="meta-llama/Meta-Llama-3-8B-Instruct",
             model_kwargs={"torch_dtype": torch.bfloat16},
             device="cuda",
         )
@@ -102,7 +101,7 @@ class PDFChatBot:
         temp = 0.1
         outputs = self.pipe(
             prompt,
-            max_new_tokens=524288,
+            max_new_tokens=1024,
             do_sample=True,
             temperature=temp,
             top_p=0.9,
@@ -114,7 +113,6 @@ class PDFChatBot:
     def process_file(self, file):
         """
         Process the uploaded PDF file and initialize necessary components: Tokenizer, VectorDB and LLM.
-
         Parameters:
             file (FileStorage): The uploaded PDF file.
         """