Update agent.py

agent.py

@@ -9,9 +9,21 @@ If you are asked for a string, don't use articles, neither abbreviations (e.g. f
 If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
 """
 
+from transformers import pipeline
+
+SYSTEM_PROMPT = """..."""  # oförändrat
+
 class GaiaAgent:
     def __init__(self, model_id="tiiuae/falcon-rw-1b"):
-        self.generator = pipeline("text-generation", model=model_id, tokenizer=model_id, max_new_tokens=512)
+        self.generator = pipeline(
+            "text-generation",
+            model=model_id,
+            tokenizer=model_id,
+            max_new_tokens=512,
+            device="cpu",  # viktigt!
+            return_full_text=False
+        )
+
 
     def __call__(self, question: str) -> tuple[str, str]:
         try: