Update src/generation/llm.py

src/generation/llm.py

@@ -99,10 +99,14 @@ 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
+            # 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
         except Exception as e: