Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import json | |
| from transformers import AutoTokenizer, AutoConfig, AutoModel | |
| # ---------------------------- | |
| # Configuration / Meta | |
| # ---------------------------- | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| MAX_LENGTH = 256 | |
| WEIGHTS_PATH = "best_model.pth" # your model file | |
| TOKENIZER_PATH = "tokenizer" | |
| META_PATH = "meta.json" | |
| # Load meta.json | |
| with open(META_PATH, "r") as f: | |
| meta = json.load(f) | |
| TARGET_COLS = meta.get("target_cols", ["anger","fear","joy","sadness","surprise"]) | |
| MAX_LENGTH = meta.get("max_length", MAX_LENGTH) | |
| BACKBONE = meta.get("backbone", "microsoft/deberta-v3-base") | |
| # ---------------------------- | |
| # Load tokenizer | |
| # ---------------------------- | |
| tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH) | |
| # ---------------------------- | |
| # Define model class | |
| # ---------------------------- | |
| class EmotionClassifier(nn.Module): | |
| def __init__(self, model_name, num_labels=5, dropout_rate=0.3): | |
| super().__init__() | |
| self.config = AutoConfig.from_pretrained(model_name) | |
| self.transformer = AutoModel.from_pretrained(model_name) | |
| self.dropout = nn.Dropout(dropout_rate) | |
| self.classifier = nn.Linear(self.config.hidden_size, num_labels) | |
| def forward(self, input_ids, attention_mask): | |
| outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask) | |
| cls = outputs.last_hidden_state[:, 0, :] | |
| cls = self.dropout(cls) | |
| logits = self.classifier(cls) | |
| return logits | |
| # ---------------------------- | |
| # Load model | |
| # ---------------------------- | |
| model = EmotionClassifier(BACKBONE, num_labels=len(TARGET_COLS)) | |
| state = torch.load(WEIGHTS_PATH, map_location=DEVICE) | |
| # Strip 'module.' prefix if present | |
| new_state = {k.replace("module.", ""): v for k,v in state.items()} | |
| model.load_state_dict(new_state, strict=False) | |
| model.to(DEVICE) | |
| model.eval() | |
| # ---------------------------- | |
| # Prediction function | |
| # ---------------------------- | |
| def predict_text(text, threshold=0.5): | |
| inputs = tokenizer( | |
| text, | |
| truncation=True, | |
| padding=True, | |
| max_length=MAX_LENGTH, | |
| return_tensors="pt" | |
| ) | |
| inputs = {k:v.to(DEVICE) for k,v in inputs.items()} | |
| with torch.no_grad(): | |
| logits = model(**inputs) | |
| probs = torch.sigmoid(logits).cpu().numpy()[0] | |
| # Percentages | |
| percentages = {label: round(float(p)*100,2) for label,p in zip(TARGET_COLS, probs)} | |
| # Binary predictions | |
| binary = {label: int(p>threshold) for label,p in zip(TARGET_COLS, probs)} | |
| # Sort probabilities descending for display | |
| percentages_sorted = dict(sorted(percentages.items(), key=lambda x: x[1], reverse=True)) | |
| return percentages_sorted, binary | |
| # ---------------------------- | |
| # Gradio Interface | |
| # ---------------------------- | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Multi-label Emotion Classifier") | |
| txt = gr.Textbox(lines=4, placeholder="Enter text here...", label="Input Text") | |
| thresh = gr.Slider(minimum=0.1, maximum=0.9, value=0.5, step=0.01, label="Threshold") | |
| btn = gr.Button("Predict") | |
| probs_out = gr.Label(label="Emotion Probabilities (%)") | |
| binary_out = gr.JSON(label="Binary Predictions (0/1)") | |
| btn.click(predict_text, inputs=[txt, thresh], outputs=[probs_out, binary_out]) | |
| if __name__ == "__main__": | |
| demo.launch() | |