Update src/generation/llm.py

src/generation/llm.py

@@ -99,10 +99,12 @@ class LLM:
             )
             # Decode the generated tokens
             response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Strip any system/user metadata
+            response = response.replace("system", "").replace("user", "").replace("assistant", "")
+            # Remove any extra whitespace or unwanted tokens
+            response = " ".join(response.split()).strip()
             print("Response generated successfully!")
-            return response.strip()
-        except Exception as e:
-            raise RuntimeError(f"Failed to generate response: {str(e)}")
+            return response
     
     def format_query_prompt(self, query: str, cuisines: List[str], dishes: List[str], price_ranges: List[str]) -> str:
         """