Update app.py

app.py

@@ -23,21 +23,14 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 
 class BasicAgent:
-    def __init__(self, model="mistral-7b-instruct-v0.1"):  # Smaller model recommended
-        self.tokenizer = AutoTokenizer.from_pretrained(model)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model,
-            device_map="auto",
-            torch_dtype=torch.float32,  # Explicitly use float32 for CPU
-            low_cpu_mem_usage=True      # Reduces memory spikes
-        )
-        print(f"Initialized on device: {self.model.device}")
-
-    def __call__(self, question: str, max_tokens: int = 100) -> str:
-        inputs = self.tokenizer(question, return_tensors="pt").to(self.model.device)
-        with torch.no_grad():  # Reduces memory usage
-            outputs = self.model.generate(**inputs, max_new_tokens=max_tokens)
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+    def __init__(self, model="google/gemma-7b"):
+        self.pipe = pipeline("text-generation", model=model)
+        print("BasicAgent initialized.")
+
+    def __call__(self, question: str) -> str:
+        print(f"Question: {question[:50]}...")
+        response = self.pipe(question, max_new_tokens=100)
+        return response[0]['generated_text']
             
     def wikipedia_search(self, query: str) -> str:
         """Get Wikipedia summary"""