Update app.py

app.py

@@ -9,7 +9,8 @@ import wikipediaapi
 import pandas as pd
 from transformers import pipeline  # or HfAgent if you want the higher-level agent
 from huggingface_hub import InferenceClient  # Updated import
-
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 load_dotenv()
 
@@ -20,22 +21,18 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
 
+
+
 class BasicAgent:
-    def __init__(self, model: str = "google/gemma-7b"):
-        self.generator = pipeline(
-            "text-generation",
-            model=model,
-            device_map="auto"  # Uses GPU if available
-        )
-    
+    def __init__(self, model="google/gemma-7b"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model)
+        self.model = AutoModelForCausalLM.from_pretrained(model)
+        print("BasicAgent initialized with AutoModel")
+
     def __call__(self, question: str) -> str:
-        outputs = self.generator(
-            question,
-            max_new_tokens=100,
-            do_sample=True,
-            temperature=0.7
-        )
-        return outputs[0]['generated_text']
+        inputs = self.tokenizer(question, return_tensors="pt")
+        outputs = self.model.generate(**inputs, max_new_tokens=100)
+        return self.tokenizer.decode(outputs[0])
             
     def wikipedia_search(self, query: str) -> str:
         """Get Wikipedia summary"""