Spaces:
Running
Running
| import gradio as gr | |
| from transformers import pipeline | |
| from smolagents import Tool | |
| class SimpleSentimentTool(Tool): | |
| name = "sentiment_analysis" | |
| 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": "oliverguhr/german-sentiment-bert", | |
| "nullable": True | |
| } | |
| } | |
| # Use a standard authorized type | |
| output_type = "string" | |
| # Available sentiment analysis models | |
| models = { | |
| "multilingual": "nlptown/bert-base-multilingual-uncased-sentiment", | |
| "deberta": "microsoft/deberta-xlarge-mnli", | |
| "distilbert": "distilbert-base-uncased-finetuned-sst-2-english", | |
| "mobilebert": "lordtt13/emo-mobilebert", | |
| "reviews": "juliensimon/reviews-sentiment-analysis", | |
| "sbc": "sbcBI/sentiment_analysis_model", | |
| "german": "oliverguhr/german-sentiment-bert" | |
| } | |
| 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 | |
| self._classifiers = {} | |
| # 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: | |
| 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 forward(self, text: str, model_key="oliverguhr/german-sentiment-bert"): | |
| """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']) | |
| # Convert to JSON string for output | |
| import json | |
| return json.dumps(result, indent=2) | |
| except Exception as e: | |
| print(f"Error in sentiment analysis: {str(e)}") | |
| return json.dumps({"error": str(e)}, indent=2) |