Update sentiment_analysis.py

sentiment_analysis.py

@@ -2,17 +2,23 @@ import gradio as gr
 from transformers import pipeline
 from smolagents import Tool
 
-class SentimentAnalysisTool(Tool):
+class SimpleSentimentTool(Tool):
     name = "sentiment_analysis"
-    description = "This tool analyses the sentiment of a given text."
+    description = "This tool analyzes the sentiment of a given text."
 
     inputs = {
         "text": {
             "type": "string",
             "description": "The text to analyze for sentiment"
+        },
+        "model_key": {
+            "type": "string",
+            "description": "The model to use for sentiment analysis",
+            "default": None
         }
     }
-    output_type = "json"
+    # Use a standard authorized type
+    output_type = "dict[str, float]"  
     
     # Available sentiment analysis models
     models = {
@@ -25,52 +31,65 @@ class SentimentAnalysisTool(Tool):
         "german": "oliverguhr/german-sentiment-bert"
     }
     
-    def __init__(self, default_model="distilbert"):
-        """Initialize with a default model."""
+    def __init__(self, default_model="distilbert", preload=False):
+        """Initialize with a default model.
+        
+        Args:
+            default_model: The default model to use if no model is specified
+            preload: Whether to preload the default model at initialization
+        """
         super().__init__()
         self.default_model = default_model
-        # Pre-load the default model to speed up first inference
         self._classifiers = {}
-        self._get_classifier(self.models[default_model])
-    
-    def forward(self, text: str): 
-        """Process input text and return sentiment predictions."""
-        return self.predict(text)
         
-    def _parse_output(self, output_json):  
-        """Parse model output into a dictionary of scores by label."""
-        result = {}
-        for i in range(len(output_json[0])):
-            label = output_json[0][i]['label']
-            score = output_json[0][i]['score']
-            result[label] = score
-        return result
+        # Optionally preload the default model
+        if preload:
+            try:
+                self._get_classifier(self.models[default_model])
+            except Exception as e:
+                print(f"Warning: Failed to preload model: {str(e)}")
     
     def _get_classifier(self, model_id):
         """Get or create a classifier for the given model ID."""
         if model_id not in self._classifiers:
-            self._classifiers[model_id] = pipeline(
-                "text-classification", 
-                model=model_id, 
-                top_k=None  # This replaces return_all_scores=True
-            )
+            try:
+                print(f"Loading model: {model_id}")
+                self._classifiers[model_id] = pipeline(
+                    "text-classification", 
+                    model=model_id, 
+                    top_k=None  # Return all scores
+                )
+            except Exception as e:
+                print(f"Error loading model {model_id}: {str(e)}")
+                # Fall back to distilbert if available
+                if model_id != self.models["distilbert"]:
+                    print("Falling back to distilbert model...")
+                    return self._get_classifier(self.models["distilbert"])
+                else:
+                    # Last resort - if even distilbert fails
+                    print("Critical error: Could not load default model")
+                    raise RuntimeError(f"Failed to load any sentiment model: {str(e)}")
         return self._classifiers[model_id]
     
-    def predict(self, text, model_key=None):
-        """Make predictions using the specified or default model."""
-        model_id = self.models[model_key] if model_key in self.models else self.models[self.default_model]
-        classifier = self._get_classifier(model_id)
-        
-        prediction = classifier(text)
-        return self._parse_output(prediction)
-
-# For standalone testing
-if __name__ == "__main__":
-    # Create an instance of the SentimentAnalysisTool class
-    sentiment_analysis_tool = SentimentAnalysisTool()
-    
-    # Test with a sample text
-    test_text = "I really enjoyed this product. It exceeded my expectations!"
-    result = sentiment_analysis_tool(test_text)
-    print(f"Input: {test_text}")
-    print(f"Result: {result}")
+    def forward(self, text: str, model_key=None):
+        """Process input text and return sentiment predictions."""
+        try:
+            # Determine which model to use
+            model_key = model_key or self.default_model
+            model_id = self.models.get(model_key, self.models[self.default_model])
+            
+            # Get the classifier
+            classifier = self._get_classifier(model_id)
+            
+            # Get predictions
+            prediction = classifier(text)
+            
+            # Format as a dictionary
+            result = {}
+            for item in prediction[0]:
+                result[item['label']] = float(item['score'])
+            
+            return result
+        except Exception as e:
+            print(f"Error in sentiment analysis: {str(e)}")
+            return {"error": str(e)}