Update src/generation/llm.py

src/generation/llm.py

@@ -99,16 +99,8 @@ class LLM:
             )
             # Decode the generated tokens
             response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            # Remove system/user/assistant metadata and unwanted tokens
-            for token in ["system", "user", "assistant", "You are", "Bạn là", "Truy vấn", "Danh sách nhà hàng", "Phản hồi"]:
-                response = response.replace(token, "")
-            # Remove extra whitespace and clean up
-            response = " ".join(response.split()).strip()
-            # Ensure the response starts cleanly (remove any leading colons or brackets)
-            if response.startswith(":") or response.startswith("["):
-                response = response[1:].strip()
             print("Response generated successfully!")
-            return response
+            return response.strip()
         except Exception as e:
             raise RuntimeError(f"Failed to generate response: {str(e)}")