Update app.py

app.py

@@ -3,8 +3,8 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
-from transformers import BartTokenizer, BartForConditionalGeneration
 import torch
+from transformers import BartTokenizer, BartForConditionalGeneration
 
 from smolagents import ToolCallingAgent
 from audio_transcriber import AudioTranscriptionTool
@@ -39,43 +39,26 @@ class LocalBartModel:
         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 or attention_mask is None:
-            raise ValueError("input_ids and attention_mask are required in inputs dict")
-
-        input_ids = input_ids.to(self.device)
-        attention_mask = attention_mask.to(self.device)
+    def __call__(self, prompt: str) -> str:
+        if not isinstance(prompt, str):
+            raise ValueError(f"LocalBartModel expects a string prompt, got {type(prompt)}")
 
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
         with torch.no_grad():
             outputs = self.model.generate(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                **generate_kwargs
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                max_length=100,
+                num_beams=5,
+                early_stopping=True,
             )
-        return outputs
-
-    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,
-            num_beams=5,
-            early_stopping=True
-        )
-        output_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        output_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         return output_text.strip()
 
 class GaiaAgent:
     def __init__(self):
         print("Gaia Agent Initialized")
+
         self.model = LocalBartModel()
 
         self.tools = [
@@ -91,19 +74,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
+                # Try to find assistant response content in list
                 for item in reversed(result):
                     if isinstance(item, dict) and item.get("role") == "assistant" and "content" in item:
                         return item["content"].strip()
@@ -244,3 +227,4 @@ if __name__ == "__main__":
 
 
 
+