Update app.py

app.py

@@ -10,7 +10,7 @@ class RAGChatbot:
                  model_name="facebook/opt-350m", 
                  embedding_model="all-MiniLM-L6-v2"):
         # Initialize tokenizer and model
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
         # self.bnb_config = BitsAndBytesConfig(
         #                 load_in_8bit=True,   # Enable 8-bit loading
         #                 llm_int8_threshold=6.0,  # Threshold for mixed-precision computation
@@ -51,7 +51,7 @@ class RAGChatbot:
             self.documents.extend(chunks)
         
         # Generate embeddings
-        self.embeddings = self.embedding_model.encode(self.documents)
+        self.embeddings = self.embedding_model.encode(self.documents, batch_size=32, show_progress_bar=True)
         return f"Loaded {len(self.documents)} text chunks from {len(file_paths)} files"
 
     def retrieve_relevant_context(self, query, top_k=3):
@@ -71,11 +71,12 @@ class RAGChatbot:
         return " ".join([self.documents[i] for i in top_indices])
 
     def generate_response(self, query, context):
-        # Construct prompt with context
-        full_prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
+        # Construct prompt with 
+        truncated_context = " ".join(context.split()[:100]) 
+        full_prompt = f"Context: {truncated_context}\n\nQuestion: {query}\n\nAnswer:"
         
         # Generate response
-        inputs = self.tokenizer(full_prompt, return_tensors="pt").to(self.model.device)
+        inputs = self.tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True).to(self.model.device)
         outputs = self.model.generate(**inputs, max_new_tokens=150)
         response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)