Update app.py

app.py

@@ -3,8 +3,13 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-from google import genai
-from google.genai import types
+# from google import genai
+# from google.genai import types
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from smolagents.agents import ReActAgent
+from smolagents.tools import tool
+
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
@@ -20,49 +25,83 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 #         print(f"Agent returning fixed answer: {fixed_answer}")
 #         return fixed_answer
 
-class BasicAgent:
+# class BasicAgent:
+#     def __init__(self):
+#         print("CustomAgent (using Gemini 2.0) initialized.")
+
+#         # Set the environment variable (important for Hugging Face Spaces)
+#         api_key = os.environ.get("GEMINI_API_KEY")
+#         if not api_key:
+#             raise ValueError("GEMINI_API_KEY not found in environment variables.")
+
+#         os.environ["GOOGLE_API_KEY"] = api_key  # Needed for google-genai Client
+
+#         # Initialize the Gemini client
+#         self.client = genai.Client()
+
+#         # Set model ID (you can choose another if needed)
+#         self.model_id = "gemini-2.0-flash-exp"
+
+#         # (Optional) Define generation config
+#         self.generation_config = types.GenerateContentConfig(
+#             temperature=0.4,
+#             top_p=0.95,
+#             top_k=20,
+#             candidate_count=1,
+#             seed=5,
+#             presence_penalty=0.0,
+#             frequency_penalty=0.0,
+#         )
+
+#     def __call__(self, question: str) -> str:
+#         print(f"Agent received question (first 50 chars): {question[:50]}...")
+
+#         try:
+#             response = self.client.models.generate_content(
+#                 model=self.model_id,
+#                 contents=f"Answer the following question clearly and concisely: {question}",
+#                 config=self.generation_config
+#             )
+#             answer = response.text.strip()
+#             print(f"Agent returning answer (first 100 chars): {answer[:100]}")
+#             return answer
+#         except Exception as e:
+#             print(f"Error during Gemini API call: {str(e)}")
+#             return f"Error: {str(e)}"
+
+class BasicAgent(ReActAgent):
     def __init__(self):
-        print("CustomAgent (using Gemini 2.0) initialized.")
-
-        # Set the environment variable (important for Hugging Face Spaces)
-        api_key = os.environ.get("GEMINI_API_KEY")
-        if not api_key:
-            raise ValueError("GEMINI_API_KEY not found in environment variables.")
-
-        os.environ["GOOGLE_API_KEY"] = api_key  # Needed for google-genai Client
-
-        # Initialize the Gemini client
-        self.client = genai.Client()
-
-        # Set model ID (you can choose another if needed)
-        self.model_id = "gemini-2.0-flash-exp"
-
-        # (Optional) Define generation config
-        self.generation_config = types.GenerateContentConfig(
-            temperature=0.4,
-            top_p=0.95,
-            top_k=20,
-            candidate_count=1,
-            seed=5,
-            presence_penalty=0.0,
-            frequency_penalty=0.0,
+        print("BasicAgent using local LLM initialized.")
+
+        # Load a small model from Hugging Face
+        model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # You can pick another lightweight model
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto"  # Automatically choose GPU/CPU
         )
 
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        super().__init__(tools=[])  # No tools for now
 
-        try:
-            response = self.client.models.generate_content(
-                model=self.model_id,
-                contents=f"Answer the following question clearly and concisely: {question}",
-                config=self.generation_config
+    def call(self, task: str) -> str:
+        """Core method for answering a task."""
+        prompt = f"Answer the following question concisely:\n\n{task}\n\nAnswer:"
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=200,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                top_k=50,
             )
-            answer = response.text.strip()
-            print(f"Agent returning answer (first 100 chars): {answer[:100]}")
-            return answer
-        except Exception as e:
-            print(f"Error during Gemini API call: {str(e)}")
-            return f"Error: {str(e)}"
+        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Extract only the answer part
+        return answer.split("Answer:")[-1].strip()
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """