Update app.py

app.py

@@ -3,13 +3,13 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
+from transformers import BartTokenizer, BartForConditionalGeneration
 import torch
-from transformers import BartForConditionalGeneration, BartTokenizer
 
+from smolagents import ToolCallingAgent
 from audio_transcriber import AudioTranscriptionTool
 from image_analyzer import ImageAnalysisTool
 from wikipedia_searcher import WikipediaSearcher
-from smolagents import ToolCallingAgent
 
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
@@ -32,32 +32,38 @@ SYSTEM_PROMPT = (
 )
 
 class LocalBartModel:
-    def __init__(self, model_name="facebook/bart-base"):
-        self.tokenizer = BartTokenizer.from_pretrained(model_name)
-        self.model = BartForConditionalGeneration.from_pretrained(model_name)
+    def __init__(self):
+        self.tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
+        self.model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.to(self.device)
+        self.model.eval()
 
     def generate(self, inputs, **generate_kwargs):
+        # inputs must be dict with input_ids and attention_mask
+        if not isinstance(inputs, dict):
+            raise ValueError(f"Expected dict input but got {type(inputs)}")
         input_ids = inputs.get("input_ids")
         attention_mask = inputs.get("attention_mask")
-
-        if input_ids is None:
-            raise ValueError("input_ids missing from tokenizer output")
+        if input_ids is None or attention_mask is None:
+            raise ValueError("input_ids and attention_mask are required in inputs dict")
 
         input_ids = input_ids.to(self.device)
-        if attention_mask is not None:
-            attention_mask = attention_mask.to(self.device)
+        attention_mask = attention_mask.to(self.device)
 
-        return self.model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            **generate_kwargs
-        )
+        with torch.no_grad():
+            outputs = self.model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                **generate_kwargs
+            )
+        return outputs
 
-    def __call__(self, prompt: str) -> str:
-        inputs = self.tokenizer(prompt, return_tensors="pt")
+    def __call__(self, prompt):
+        if not isinstance(prompt, str):
+            raise ValueError(f"LocalBartModel expects a string prompt, got {type(prompt)}")
 
+        inputs = self.tokenizer(prompt, return_tensors="pt")
         output_ids = self.generate(
             inputs,
             max_length=100,
@@ -71,11 +77,13 @@ class GaiaAgent:
     def __init__(self):
         print("Gaia Agent Initialized")
         self.model = LocalBartModel()
+
         self.tools = [
             AudioTranscriptionTool(),
             ImageAnalysisTool(),
             WikipediaSearcher()
         ]
+
         self.agent = ToolCallingAgent(
             tools=self.tools,
             model=self.model
@@ -83,18 +91,19 @@ class GaiaAgent:
 
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-
         full_prompt = f"{SYSTEM_PROMPT}\nQUESTION:\n{question}"
 
         try:
             result = self.agent.run(full_prompt)
             print(f"Raw result from agent: {result}")
 
+            # Handle different result types robustly
             if isinstance(result, dict) and "answer" in result:
                 return str(result["answer"]).strip()
             elif isinstance(result, str):
                 return result.strip()
             elif isinstance(result, list):
+                # Try to extract assistant content from list
                 for item in reversed(result):
                     if isinstance(item, dict) and item.get("role") == "assistant" and "content" in item:
                         return item["content"].strip()