Spaces:
Running
Running
| from typing import Any, Dict, Union, Tuple | |
| import gradio as gr | |
| from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer | |
| import logging | |
| from .base import BaseModel | |
| logger = logging.getLogger(__name__) | |
| class TextClassificationModel(BaseModel): | |
| """Lightweight text classification model using tiny BERT.""" | |
| def __init__(self): | |
| super().__init__( | |
| name="Lightweight Text Classifier", | |
| description="Fast text classification using a tiny BERT model (4.4MB)" | |
| ) | |
| self.model_name = "prajjwal1/bert-tiny" | |
| self._model = None | |
| def load_model(self) -> None: | |
| """Load the classification model.""" | |
| try: | |
| logger.info(f"Loading model: {self.model_name}") | |
| # Initialize model with binary classification | |
| model = AutoModelForSequenceClassification.from_pretrained( | |
| self.model_name, | |
| num_labels=2 | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained(self.model_name) | |
| self._model = pipeline( | |
| "text-classification", | |
| model=model, | |
| tokenizer=tokenizer, | |
| device=-1 # CPU, use device=0 for GPU | |
| ) | |
| # Log model size | |
| model_size_mb = sum(p.numel() * p.element_size() for p in model.parameters()) / (1024 * 1024) | |
| logger.info(f"Model loaded successfully. Size: {model_size_mb:.2f} MB") | |
| except Exception as e: | |
| logger.error(f"Error loading model: {str(e)}") | |
| raise | |
| async def predict(self, text: str) -> Dict[str, Union[str, float]]: | |
| """Make a prediction using the model.""" | |
| try: | |
| if self._model is None: | |
| self.load_model() | |
| logger.info(f"Processing text: {text[:50]}...") | |
| result = self._model(text)[0] | |
| # Map raw labels to sentiment | |
| label_map = { | |
| "LABEL_0": "NEGATIVE", | |
| "LABEL_1": "POSITIVE" | |
| } | |
| prediction = { | |
| "label": label_map.get(result["label"], result["label"]), | |
| "confidence": float(result["score"]) | |
| } | |
| logger.info(f"Prediction result: {prediction}") | |
| return prediction | |
| except Exception as e: | |
| logger.error(f"Prediction error: {str(e)}") | |
| raise | |
| async def predict_for_interface(self, text: str) -> Tuple[str, float]: | |
| """Make a prediction and return it in a format suitable for the Gradio interface.""" | |
| result = await self.predict(text) | |
| return result["label"], result["confidence"] | |
| def create_interface(self) -> gr.Interface: | |
| """Create a Gradio interface for text classification.""" | |
| if self._model is None: | |
| self.load_model() | |
| examples = [ | |
| ["This movie was fantastic! I really enjoyed it."], | |
| ["The service was terrible and the food was cold."], | |
| ["It was an okay experience, nothing special."], | |
| ["The weather is nice today!"], | |
| ["I'm feeling sick and tired."] | |
| ] | |
| return gr.Interface( | |
| fn=self.predict_for_interface, # Use the interface-specific prediction function | |
| inputs=gr.Textbox( | |
| lines=3, | |
| placeholder="Enter text to classify...", | |
| label="Input Text" | |
| ), | |
| outputs=[ | |
| gr.Label(label="Sentiment"), | |
| gr.Number(label="Confidence", precision=4) | |
| ], | |
| title=self.name, | |
| description=self.description + "\n\nThis model is also available via API!", | |
| examples=examples, | |
| api_name="predict" | |
| ) |