Update app.py

app.py

@@ -14,20 +14,14 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 class BasicAgent:
     def __init__(self):
-        print("Mistral Agent loading on CPU...")
-
-        model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+        print("Falcon-RW-1B Agent initialized.")
+        model_id = "tiiuae/falcon-rw-1b"
 
         # Load tokenizer and model
         self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            device_map="auto",   # Will default to CPU
-            low_cpu_mem_usage=True,  # Helps a bit
-            torch_dtype="auto"
-        )
+        self.model = AutoModelForCausalLM.from_pretrained(model_id)
 
-        # Create pipeline (CPU-only)
+        # Create a CPU-based pipeline
         self.pipeline = pipeline(
             "text-generation",
             model=self.model,
@@ -39,23 +33,20 @@ class BasicAgent:
         print(f"Agent received question: {question[:50]}...")
 
         try:
-            # Format with instruction template
-            prompt = f"<s>[INST] {question.strip()} [/INST]"
-
+            prompt = f"Question: {question.strip()}\nAnswer:"
             output = self.pipeline(
                 prompt,
-                max_new_tokens=256,
-                do_sample=True,
+                max_new_tokens=128,
                 temperature=0.7,
-                top_p=0.95
+                top_p=0.9
             )
 
             full_response = output[0]["generated_text"]
-            answer = full_response.split("[/INST]")[-1].strip()
-
+            answer = full_response.split("Answer:")[-1].strip()
             return answer
+
         except Exception as e:
-            print(f"❌ Mistral error: {e}")
+            print(f"❌ Falcon error: {e}")
             return f"❌ Model Error: {str(e)}"
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):