Update main.py

main.py

@@ -123,46 +123,34 @@ User query: {query}
         try:
             context_chunks = self.retrieve_relevant_chunks(query)
             prompt = self.format_prompt(query, context_chunks)
-            input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
             
-            max_length = input_ids.shape[1] + 512
-            generated_text = ""
-
+            input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
+    
             with torch.no_grad():
-                generated = input_ids
-                
-                while generated.shape[1] < max_length:
-                    outputs = self.model.forward(
-                        input_ids=generated,
-                        max_new_tokens=64,
-                        do_sample=True,
-                        temperature=0.7,
-                        top_p=0.8,
-                        repetition_penalty=1.05,
-                        pad_token_id=self.tokenizer.pad_token_id,
-                        eos_token_id=self.tokenizer.eos_token_id,
-                    )
-                    
-                    next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1).unsqueeze(0)
-                    
-                    if next_token[0, 0].item() == self.tokenizer.eos_token_id:
-                        break
-                        
-                    generated = torch.cat([generated, next_token.T], dim=1)
-                    new_text = self.tokenizer.decode(next_token[0], skip_special_tokens=True)
-                    
-                    if new_text:
-                        generated_text += new_text
-            
+                outputs = self.model.generate(   # Fixed: was model.forward()
+                    input_ids=input_ids,
+                    max_new_tokens=512,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.8,
+                    repetition_penalty=1.05,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                )
+    
+            # Decode only the newly generated tokens (strip the prompt)
+            new_tokens = outputs[0][input_ids.shape[1]:]
+            generated_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
+    
             return {
                 "Response_text": generated_text,
                 "context_items": context_chunks
             }
-            
+    
         except Exception as e:
             print(f"Error generating response: {str(e)}")
             return {
-                "Response_text": "I apologize, but I encountered an error while processing your query. Please try again.",
+                "Response_text": "I apologize, but I encountered an error. Please try again.",
                 "context_items": []
             }