Spaces:
Running
Running
| import pandas as pd | |
| import numpy as np | |
| import torch | |
| import re | |
| import emoji | |
| import matplotlib.pyplot as plt | |
| import seaborn as sns | |
| from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix, classification_report | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from setfit import SetFitModel | |
| from transformers import AutoTokenizer, AutoModelForMaskedLM | |
| import io | |
| import tempfile | |
| import os | |
| from PIL import Image | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| HF_USERNAME = "Methni" | |
| SETFIT_REPO = f"{HF_USERNAME}/STEMO-SetFit" | |
| DATASET_REPO = f"{HF_USERNAME}/STEMO-Dataset" | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Using device: {device}") | |
| EMOTION_INFO = { | |
| 'Happy': {'emoji': '😊', 'color': '#FFD700', 'description': 'Joy, excitement, positivity'}, | |
| 'Anger': {'emoji': '😠', 'color': '#FF4444', 'description': 'Frustration, rage, irritation'}, | |
| 'Sadness': {'emoji': '😢', 'color': '#4169E1', 'description': 'Grief, disappointment, sorrow'}, | |
| 'Fear': {'emoji': '😨', 'color': '#9370DB', 'description': 'Worry, anxiety, dread'}, | |
| 'Surprise': {'emoji': '😲', 'color': '#FF8C00', 'description': 'Shock, astonishment, disbelief'}, | |
| 'Disgust': {'emoji': '🤢', 'color': '#228B22', 'description': 'Revulsion, distaste, contempt'}, | |
| } | |
| MODEL_INFO = { | |
| 'SetFit (Recommended)': { | |
| 'key': 'setfit', | |
| 'accuracy': '80.65%', | |
| 'description': 'Best overall accuracy. Recommended for most users.', | |
| }, | |
| 'Prompt-Based': { | |
| 'key': 'fewshot', | |
| 'accuracy': '58.71%', | |
| 'description': 'Works without any training. More robust to very noisy text.', | |
| }, | |
| } | |
| # PREPROCESSING | |
| def clean_text_setfit(text): | |
| if not isinstance(text, str): return "" | |
| text = re.sub(r'http\S+', '', text) | |
| text = re.sub(r'@\w+', '', text) | |
| text = re.sub(r'\s+', ' ', text).strip() | |
| text = emoji.demojize(text) | |
| return text | |
| def clean_text_fewshot(text): | |
| if not isinstance(text, str): return "" | |
| text = re.sub(r'http\S+', '', text) | |
| text = re.sub(r'@\w+', '', text) | |
| text = re.sub(r'\s+', ' ', text).strip() | |
| text = emoji.demojize(text, delimiters=(" ", " ")) | |
| return text | |
| def detect_language_stats(text): | |
| text_no_emoji = re.sub(r':[a-z_]+:', '', text) | |
| text_no_emoji = re.sub(r'\b[a-z]+_[a-z_]+\b', '', text_no_emoji) | |
| sinhala = len(re.findall(r'[\u0D80-\u0DFF]', text_no_emoji)) | |
| tamil = len(re.findall(r'[\u0B80-\u0BFF]', text_no_emoji)) | |
| english = len(re.findall(r'[a-zA-Z]', text_no_emoji)) | |
| total = sinhala + tamil + english | |
| if total == 0: | |
| return {'sinhala': 0, 'tamil': 0, 'english': 0, 'is_code_mixed': False} | |
| return { | |
| 'sinhala': sinhala / total, | |
| 'tamil': tamil / total, | |
| 'english': english / total, | |
| 'is_code_mixed': sinhala > 0 and tamil > 0 | |
| } | |
| # FEW-SHOT COMPONENTS | |
| class SmartExampleSelector: | |
| def __init__(self, support_df): | |
| self.support_df = support_df.reset_index(drop=True) | |
| self.vectorizer = TfidfVectorizer(analyzer='char', ngram_range=(3, 5), min_df=2) | |
| cleaned = [clean_text_fewshot(t) for t in self.support_df['text']] | |
| self.support_vecs = self.vectorizer.fit_transform(cleaned) | |
| print(f" Selector ready with {len(self.support_df)} examples") | |
| def get_k_similar(self, query_text, k=3): | |
| query_vec = self.vectorizer.transform([query_text]) | |
| sim_scores = cosine_similarity(query_vec, self.support_vecs).flatten() | |
| top_indices = sim_scores.argsort()[-k:][::-1] | |
| return self.support_df.iloc[top_indices] | |
| class FewShotClassifier: | |
| def __init__(self): | |
| print("Loading few-shot model...") | |
| self.tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base") | |
| self.model = AutoModelForMaskedLM.from_pretrained("xlm-roberta-base").to(device) | |
| self.model.eval() | |
| self.label_map = { | |
| 'Happy': 'happy', 'Anger': 'mad', 'Sadness': 'sad', | |
| 'Fear': 'fear', 'Surprise': 'shock', 'Disgust': 'gross' | |
| } | |
| self.logit_bias = {l: 0.0 for l in self.label_map} | |
| self.labels = list(self.label_map.keys()) | |
| self.verbalizer_ids = [ | |
| self.tokenizer.encode(self.label_map[l], add_special_tokens=False)[0] | |
| for l in self.labels | |
| ] | |
| print("Few-shot model loaded") | |
| def get_logits(self, prompt): | |
| inputs = self.tokenizer(prompt, return_tensors="pt", | |
| truncation=True, max_length=512).to(device) | |
| mask_idx = (inputs.input_ids == self.tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0] | |
| if len(mask_idx) == 0: | |
| return torch.zeros(len(self.verbalizer_ids)) | |
| with torch.no_grad(): | |
| outputs = self.model(**inputs) | |
| mask_logits = outputs.logits[0, mask_idx[0], :] | |
| return torch.tensor([mask_logits[vid].item() for vid in self.verbalizer_ids]) | |
| def predict(self, text, examples): | |
| prompt = "" | |
| for _, row in examples.iterrows(): | |
| prompt += f"Tweet: {row['text']} \nEmotion: {self.label_map[row['label']]}\n\n" | |
| prompt += f"Tweet: {text} \nEmotion: {self.tokenizer.mask_token}" | |
| raw = self.get_logits(prompt) | |
| null_prompt = f"Tweet: [N/A] \nEmotion: {self.tokenizer.mask_token}" | |
| bias = self.get_logits(null_prompt) | |
| scores = [(raw[i] - bias[i]) + self.logit_bias[l] for i, l in enumerate(self.labels)] | |
| probs = torch.softmax(torch.tensor(scores), dim=0).cpu().numpy() | |
| return self.labels[int(np.argmax(scores))], probs | |
| # MODEL MANAGER | |
| class ModelManager: | |
| def __init__(self): | |
| self.label_names = ['Happy', 'Anger', 'Sadness', 'Fear', 'Surprise', 'Disgust'] | |
| self.setfit_model = None | |
| self.fewshot_classifier = None | |
| self.fewshot_selector = None | |
| def load_all_models(self): | |
| print("=" * 60) | |
| print("LOADING MODELS FROM HUGGING FACE") | |
| print("=" * 60) | |
| # 1. SetFit | |
| try: | |
| print("\n1. Loading SetFit...") | |
| self.setfit_model = SetFitModel.from_pretrained(SETFIT_REPO) | |
| self.setfit_model.to(device) | |
| print(" ✓ SetFit loaded") | |
| except Exception as e: | |
| print(f" ✗ {e}") | |
| # 2. Few-shot | |
| try: | |
| print("\n2. Loading Few-Shot components...") | |
| train_path = hf_hub_download( | |
| repo_id=DATASET_REPO, | |
| filename="STEMO_Train_Raw.xlsx", | |
| repo_type="dataset" | |
| ) | |
| train_df = pd.read_excel(train_path) | |
| self.fewshot_selector = SmartExampleSelector(train_df) | |
| self.fewshot_classifier = FewShotClassifier() | |
| print(" ✓ Few-shot loaded") | |
| except Exception as e: | |
| print(f" ✗ {e}") | |
| print("\n" + "=" * 60) | |
| print("ALL MODELS LOADED — STEMO READY") | |
| print("=" * 60) | |
| def predict_setfit(self, text): | |
| if self.setfit_model is None: | |
| return {'error': 'SetFit model not loaded'} | |
| processed = clean_text_setfit(text) | |
| lang = detect_language_stats(processed) | |
| try: | |
| predictions = self.setfit_model.predict([processed]) | |
| probs = self.setfit_model.predict_proba([processed])[0] | |
| prediction = predictions[0] | |
| confidence = float(probs.max()) | |
| emotion = (self.label_names[prediction] | |
| if isinstance(prediction, (int, np.integer)) else prediction) | |
| result = { | |
| 'model': 'SetFit (Recommended)', 'emotion': emotion, | |
| 'confidence': confidence, | |
| 'all_scores': {self.label_names[i]: float(probs[i]) for i in range(len(probs))}, | |
| 'lang': lang, | |
| } | |
| if confidence < 0.5: result['warning'] = True | |
| return result | |
| except Exception as e: | |
| return {'error': str(e)} | |
| def predict_fewshot(self, text): | |
| if self.fewshot_classifier is None: | |
| return {'error': 'Few-shot model not loaded'} | |
| processed = clean_text_fewshot(text) | |
| lang = detect_language_stats(processed) | |
| try: | |
| examples = self.fewshot_selector.get_k_similar(processed, k=3) | |
| emotion, probs = self.fewshot_classifier.predict(processed, examples) | |
| confidence = float(probs.max()) | |
| result = { | |
| 'model': 'Prompt-Based', 'emotion': emotion, 'confidence': confidence, | |
| 'all_scores': {self.label_names[i]: float(probs[i]) for i in range(len(probs))}, | |
| 'lang': lang, | |
| } | |
| if confidence < 0.5: result['warning'] = True | |
| return result | |
| except Exception as e: | |
| return {'error': str(e)} | |
| def predict_all(self, text): | |
| results = {} | |
| if self.setfit_model: results['SetFit (Recommended)'] = self.predict_setfit(text) | |
| if self.fewshot_classifier: results['Prompt-Based'] = self.predict_fewshot(text) | |
| return results | |
| def predict_by_key(self, text, key): | |
| if key == 'setfit': return self.predict_setfit(text) | |
| if key == 'fewshot': return self.predict_fewshot(text) | |
| return {'error': 'Unknown model'} | |
| #INITIALIZE | |
| model_manager = ModelManager() | |
| model_manager.load_all_models() | |
| # VISUALIZATION HELPERS | |
| def build_confidence_chart(all_scores): | |
| emotions = list(EMOTION_INFO.keys()) | |
| scores = [all_scores.get(e, 0) for e in emotions] | |
| colors = [EMOTION_INFO[e]['color'] for e in emotions] | |
| emojis = [EMOTION_INFO[e]['emoji'] for e in emotions] | |
| labels = [f"{emojis[i]} {emotions[i]}" for i in range(len(emotions))] | |
| fig, ax = plt.subplots(figsize=(10, 5)) | |
| fig.patch.set_facecolor('#FAFAFA') | |
| ax.set_facecolor('#FAFAFA') | |
| bars = ax.barh(labels, scores, color=colors, alpha=0.85, | |
| edgecolor='white', linewidth=1.5, height=0.6) | |
| max_idx = scores.index(max(scores)) | |
| bars[max_idx].set_edgecolor('#333333') | |
| bars[max_idx].set_linewidth(2.5) | |
| ax.set_xlabel('Confidence (higher = more certain)', fontsize=11) | |
| ax.set_title('How confident is the model about each emotion?', | |
| fontsize=13, fontweight='bold', pad=15) | |
| ax.set_xlim([0, 1.15]) | |
| ax.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f'{x:.0%}')) | |
| ax.tick_params(axis='y', labelsize=11) | |
| ax.tick_params(axis='x', labelsize=10) | |
| for i, (bar, score) in enumerate(zip(bars, scores)): | |
| ax.text(score + 0.02, bar.get_y() + bar.get_height() / 2, | |
| f'{score:.0%}', va='center', ha='left', fontsize=10, | |
| fontweight='bold' if i == max_idx else 'normal', color='#222222') | |
| ax.grid(axis='x', linestyle='--', alpha=0.4) | |
| ax.spines['top'].set_visible(False) | |
| ax.spines['right'].set_visible(False) | |
| plt.tight_layout() | |
| buf = io.BytesIO() | |
| plt.savefig(buf, format='png', dpi=120, bbox_inches='tight') | |
| buf.seek(0); img = Image.open(buf); plt.close() | |
| return img | |
| def build_comparison_chart(results): | |
| model_names, confidences, emotions = [], [], [] | |
| for name, r in results.items(): | |
| if 'error' not in r: | |
| model_names.append(name) | |
| confidences.append(r['confidence']) | |
| emotions.append(r['emotion']) | |
| if not model_names: return None | |
| fig, ax = plt.subplots(figsize=(10, 5)) | |
| fig.patch.set_facecolor('#FAFAFA') | |
| ax.set_facecolor('#FAFAFA') | |
| bar_colors = [EMOTION_INFO.get(e, {}).get('color', '#888888') for e in emotions] | |
| bars = ax.bar(model_names, confidences, color=bar_colors, alpha=0.85, | |
| edgecolor='white', linewidth=2, width=0.5) | |
| ax.axhline(y=0.5, color='#CC0000', linestyle='--', linewidth=1.5, | |
| label='Low-confidence threshold (50%)') | |
| ax.set_ylabel('Confidence Score', fontsize=11) | |
| ax.set_title('Which model is most confident — and what did each predict?', | |
| fontsize=13, fontweight='bold', pad=15) | |
| ax.set_ylim([0, 1.25]) | |
| ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f'{x:.0%}')) | |
| ax.tick_params(axis='x', labelsize=10) | |
| ax.tick_params(axis='y', labelsize=10) | |
| ax.legend(fontsize=9) | |
| for bar, conf, emo in zip(bars, confidences, emotions): | |
| info = EMOTION_INFO.get(emo, {}) | |
| em = info.get('emoji', '') | |
| ax.text(bar.get_x() + bar.get_width() / 2, conf + 0.03, | |
| f'{em} {emo}\n{conf:.0%}', | |
| ha='center', va='bottom', fontsize=10, fontweight='bold') | |
| ax.spines['top'].set_visible(False) | |
| ax.spines['right'].set_visible(False) | |
| ax.grid(axis='y', linestyle='--', alpha=0.4) | |
| plt.tight_layout() | |
| buf = io.BytesIO() | |
| plt.savefig(buf, format='png', dpi=120, bbox_inches='tight') | |
| buf.seek(0); img = Image.open(buf); plt.close() | |
| return img | |
| def build_result_card(result): | |
| if 'error' in result: | |
| return f" **Something went wrong:** {result['error']}\n\nPlease check your input and try again." | |
| emotion = result['emotion'] | |
| info = EMOTION_INFO.get(emotion, {}) | |
| em = info.get('emoji', '❓') | |
| conf = result['confidence'] | |
| lang = result.get('lang', {}) | |
| if conf >= 0.75: conf_label = "🟢 High confidence" | |
| elif conf >= 0.50: conf_label = "🟡 Moderate confidence" | |
| else: conf_label = "🔴 Low confidence — the model is uncertain about this tweet" | |
| lang_parts = [] | |
| if lang.get('sinhala', 0) > 0.05: lang_parts.append(f"Sinhala ({lang['sinhala']:.0%})") | |
| if lang.get('tamil', 0) > 0.05: lang_parts.append(f"Tamil ({lang['tamil']:.0%})") | |
| if lang.get('english', 0) > 0.05: lang_parts.append(f"English ({lang['english']:.0%})") | |
| lang_str = " + ".join(lang_parts) if lang_parts else "Not detected" | |
| mixed_str = "Yes — this is a code-mixed tweet" if lang.get('is_code_mixed') else "No" | |
| output = f"""## {em} The detected emotion is **{emotion}** | |
| > *{info.get('description', '')}* | |
| | | | | |
| |---|---| | |
| | **Confidence** | {conf:.0%} — {conf_label} | | |
| | **Languages detected** | {lang_str} | | |
| | **Code-mixed?** | {mixed_str} | | |
| --- | |
| """ | |
| if result.get('warning'): | |
| output += ( | |
| "\n> **Note:** The model is less than 50% confident about this result. " | |
| "This can happen with very short tweets, unusual spelling, or mixed scripts. " | |
| "Try rewording or using a different model.\n\n" | |
| ) | |
| return output | |
| def save_to_xlsx(df, filename="results.xlsx"): | |
| path = os.path.join(tempfile.gettempdir(), filename) | |
| df.to_excel(path, index=False) | |
| return path | |
| # ── GRADIO TAB FUNCTIONS ────────────────────────────────────── | |
| def tab_single(text, model_display_name): | |
| if not text.strip(): | |
| return ("**Please type or paste a tweet above and click Analyse.**\n\n" | |
| "You can mix Sinhala, Tamil, and English — emojis are welcome too! 😊", | |
| None, None) | |
| key = MODEL_INFO[model_display_name]['key'] | |
| result = model_manager.predict_by_key(text, key) | |
| card = build_result_card(result) | |
| chart = build_confidence_chart(result['all_scores']) if 'error' not in result else None | |
| if 'error' not in result: | |
| rows = sorted(result['all_scores'].items(), key=lambda x: -x[1]) | |
| table = pd.DataFrame([{ | |
| 'Emotion': f"{EMOTION_INFO[e]['emoji']} {e}", | |
| 'Confidence': f"{s:.0%}", | |
| 'What it means': EMOTION_INFO[e]['description'] | |
| } for e, s in rows]) | |
| else: | |
| table = None | |
| return card, table, chart | |
| def tab_compare(text): | |
| if not text.strip(): | |
| return ("**Please type or paste a tweet above and click Compare.**", None, None) | |
| results = model_manager.predict_all(text) | |
| chart = build_comparison_chart(results) | |
| output = "## Model Comparison Results\n\n" | |
| output += ("Each STEMO model was run on your tweet independently. " | |
| "The table below shows what each model predicted and how confident it was.\n\n") | |
| rows = [] | |
| for name, r in results.items(): | |
| if 'error' not in r: | |
| em = EMOTION_INFO.get(r['emotion'], {}).get('emoji', '') | |
| conf = r['confidence'] | |
| note = ("🟢 High confidence" if conf >= 0.75 else | |
| "🟡 Moderate confidence" if conf >= 0.50 else "🔴 Low confidence") | |
| rows.append({'Model': name, | |
| 'Predicted Emotion': f"{em} {r['emotion']}", | |
| 'Confidence': f"{conf:.0%}", | |
| 'Confidence Level': note}) | |
| table = pd.DataFrame(rows) if rows else None | |
| if rows: | |
| emotions_predicted = [r['emotion'] for r in results.values() if 'error' not in r] | |
| if len(set(emotions_predicted)) == 1: | |
| output += f" **All models agree: the emotion is {emotions_predicted[0]}.**\n\n" | |
| else: | |
| output += ("**The models disagree on this tweet.** " | |
| "This often happens when the tweet is ambiguous or very short. " | |
| "The **SetFit (Recommended)** result is usually the most reliable.\n\n") | |
| return output, table, chart | |
| def tab_batch(file, model_display_name, text_col): | |
| if file is None: | |
| return ("**Please upload a file to get started.**\n\n" | |
| "Your file should be a spreadsheet (.xlsx) or CSV. " | |
| "The default tweet column name is **text**.", None, None) | |
| key = MODEL_INFO[model_display_name]['key'] | |
| try: | |
| df = pd.read_excel(file.name) if file.name.endswith('.xlsx') else pd.read_csv(file.name) | |
| if text_col not in df.columns: | |
| return (f"**Column '{text_col}' was not found.**\n\n" | |
| f"Available columns: **{', '.join(df.columns)}**", None, None) | |
| results_list = [] | |
| for _, row in df.iterrows(): | |
| result = model_manager.predict_by_key(str(row[text_col]), key) | |
| entry = row.to_dict() | |
| if 'error' not in result: | |
| em = EMOTION_INFO.get(result['emotion'], {}).get('emoji', '') | |
| entry['Predicted Emotion'] = f"{em} {result['emotion']}" | |
| entry['Confidence'] = f"{result['confidence']:.0%}" | |
| entry['Confidence Level'] = ('High' if result['confidence'] >= 0.75 else | |
| 'Moderate' if result['confidence'] >= 0.50 else | |
| 'Low — review manually') | |
| else: | |
| entry['Predicted Emotion'] = 'Error' | |
| entry['Confidence'] = 'N/A' | |
| entry['Confidence Level'] = 'Error' | |
| results_list.append(entry) | |
| results_df = pd.DataFrame(results_list) | |
| path = save_to_xlsx(results_df, "STEMO_Batch_Results.xlsx") | |
| top_emotion = results_df['Predicted Emotion'].value_counts().index[0] | |
| status = (f"## Analysis Complete!\n\n" | |
| f"| | |\n|---|---|\n" | |
| f"| **Tweets analysed** | {len(results_df)} |\n" | |
| f"| **Model used** | {model_display_name} |\n" | |
| f"| **Most common emotion** | {top_emotion} |\n\n" | |
| f"Click **Download Results** to save the full analysis.") | |
| return status, results_df.head(10), path | |
| except Exception as e: | |
| return f"**Error:** {str(e)}", None, None | |
| def tab_evaluate(file, model_display_name, text_col, label_col): | |
| if file is None: | |
| return ("**Please upload a labelled test file.**\n\n" | |
| "Your file needs a **text** column and a **label** column.\n\n" | |
| "Valid labels: Happy, Anger, Sadness, Fear, Surprise, Disgust", | |
| None, None, None) | |
| key = MODEL_INFO[model_display_name]['key'] | |
| try: | |
| df = pd.read_excel(file.name) if file.name.endswith('.xlsx') else pd.read_csv(file.name) | |
| if text_col not in df.columns or label_col not in df.columns: | |
| return (f" **Column not found.** Available: {', '.join(df.columns)}", | |
| None, None, None) | |
| y_true, y_pred, confs = [], [], [] | |
| for _, row in df.iterrows(): | |
| result = model_manager.predict_by_key(str(row[text_col]), key) | |
| if 'error' not in result: | |
| y_true.append(row[label_col]) | |
| y_pred.append(result['emotion']) | |
| confs.append(result['confidence']) | |
| acc = accuracy_score(y_true, y_pred) | |
| prec, rec, f1, _ = precision_recall_fscore_support( | |
| y_true, y_pred, average='macro', zero_division=0) | |
| output = (f"## Evaluation Results — {model_display_name}\n\n" | |
| f"| Metric | Score | What it means |\n|---|---|---|\n" | |
| f"| **Accuracy** | {acc:.1%} | Out of every 100 tweets, the model got this many right |\n" | |
| f"| **Precision** | {prec:.1%} | When the model predicts an emotion, how often it is correct |\n" | |
| f"| **Recall** | {rec:.1%} | How well the model finds all tweets with each emotion |\n" | |
| f"| **F1 Score** | {f1:.1%} | Overall balance between precision and recall |\n" | |
| f"| **Avg Confidence** | {np.mean(confs):.1%} | Average certainty of predictions |\n\n" | |
| f"---\n**Detailed breakdown by emotion:**\n" | |
| f"```\n{classification_report(y_true, y_pred, zero_division=0)}\n```") | |
| metrics_df = pd.DataFrame([ | |
| {'Metric': 'Accuracy', 'Score': f"{acc:.1%}"}, | |
| {'Metric': 'Precision', 'Score': f"{prec:.1%}"}, | |
| {'Metric': 'Recall', 'Score': f"{rec:.1%}"}, | |
| {'Metric': 'F1 Score', 'Score': f"{f1:.1%}"}, | |
| {'Metric': 'Avg Confidence', 'Score': f"{np.mean(confs):.1%}"}, | |
| ]) | |
| cm = confusion_matrix(y_true, y_pred, labels=model_manager.label_names) | |
| fig, ax = plt.subplots(figsize=(10, 8)) | |
| fig.patch.set_facecolor('#FAFAFA') | |
| sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', | |
| xticklabels=[f"{EMOTION_INFO[l]['emoji']} {l}" for l in model_manager.label_names], | |
| yticklabels=[f"{EMOTION_INFO[l]['emoji']} {l}" for l in model_manager.label_names], | |
| cbar_kws={'label': 'Number of tweets'}, linewidths=0.5) | |
| ax.set_title(f'{model_display_name} — Confusion Matrix\n' | |
| 'Diagonal = correct predictions | Off-diagonal = mistakes', | |
| fontsize=13, fontweight='bold', pad=15) | |
| ax.set_ylabel('Actual Emotion', fontsize=11) | |
| ax.set_xlabel('Predicted Emotion', fontsize=11) | |
| plt.xticks(rotation=30, ha='right'); plt.yticks(rotation=0) | |
| plt.tight_layout() | |
| buf = io.BytesIO() | |
| plt.savefig(buf, format='png', dpi=120, bbox_inches='tight') | |
| buf.seek(0); img = Image.open(buf); plt.close() | |
| results_df = df[[text_col, label_col]].copy() | |
| results_df['Predicted'] = y_pred | |
| results_df['Correct?'] = results_df[label_col] == results_df['Predicted'] | |
| results_df['Confidence'] = [f"{c:.0%}" for c in confs] | |
| per_class = classification_report(y_true, y_pred, output_dict=True, zero_division=0) | |
| per_class_df = pd.DataFrame(per_class).T.reset_index().rename(columns={'index': 'Class'}) | |
| path = os.path.join(tempfile.gettempdir(), "STEMO_Evaluation_Results.xlsx") | |
| with pd.ExcelWriter(path, engine='openpyxl') as writer: | |
| results_df.to_excel(writer, sheet_name='Predictions', index=False) | |
| metrics_df.to_excel(writer, sheet_name='Overall Metrics', index=False) | |
| per_class_df.to_excel(writer, sheet_name='Per Emotion F1', index=False) | |
| return output, metrics_df, img, path | |
| except Exception as e: | |
| return f"**Error:** {str(e)}", None, None, None | |
| #GRADIO UI | |
| MODEL_CHOICES = list(MODEL_INFO.keys()) | |
| with gr.Blocks(theme=gr.themes.Soft(), title="STEMO — Emotion Classifier") as demo: | |
| gr.HTML(""" | |
| <div style="text-align:center; background:linear-gradient(135deg,#4f46e5,#7c3aed); | |
| padding:36px 20px; border-radius:16px; margin-bottom:24px;"> | |
| <h1 style="color:white; margin:0; font-size:2.4em; font-weight:700; letter-spacing:-0.5px;"> | |
| STEMO | |
| </h1> | |
| <p style="color:rgba(255,255,255,0.95); margin:10px 0 4px; font-size:1.25em; font-weight:500;"> | |
| Sinhala-Tamil Emotion Classifier | |
| </p> | |
| <p style="color:rgba(255,255,255,0.75); margin:0; font-size:1em;"> | |
| Type any tweet in Sinhala, Tamil, English, or a mix and discover its emotion instantly. | |
| </p> | |
| </div> | |
| """) | |
| gr.HTML(""" | |
| <div style="background:#f0fdf4; border:1px solid #bbf7d0; border-radius:12px; | |
| padding:16px 20px; margin-bottom:20px;"> | |
| <p style="margin:0; font-size:0.97em; color:#166534;"> | |
| <strong>How STEMO works:</strong> | |
| STEMO reads your tweet including Sinhala (සිංහල), Tamil (தமிழ்), emojis, and English | |
| and classifies it into one of six emotions: | |
| <strong>Happy | Anger | Sadness | | |
| Fear | Surprise | Disgust</strong>. | |
| </p> | |
| </div> | |
| """) | |
| with gr.Tabs(): | |
| # ── TAB 1 ───────────────────────────────────────────── | |
| with gr.Tab("Analyse a Tweet"): | |
| gr.Markdown("### Step 1 — Type or paste your tweet below\n" | |
| "You can write in Sinhala, Tamil, English, or mix freely. Emojis help! 😊") | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| t1_text = gr.Textbox(label="Your tweet", lines=3, | |
| placeholder="e.g. මං ගොඩක් සතුටුයි today! 🎉", | |
| max_lines=6) | |
| with gr.Column(scale=2): | |
| gr.HTML(""" | |
| <div style="background:#fefce8; border:1px solid #fde047; | |
| border-radius:10px; padding:12px 14px;"> | |
| <p style="margin:0 0 8px; font-weight:600; font-size:0.95em; color:#713f12;"> | |
| Try one of these example tweets: | |
| </p> | |
| <p style="margin:0; font-size:0.88em; color:#92400e; line-height:1.8;"> | |
| <em>මං ගොඩක් සතුටුයි අද 😊</em><br> | |
| <em>I'm really කනගාටුයි about this</em><br> | |
| <em>மிகவும் கோபமா இருக்கு today!</em><br> | |
| <em>OMG இது என்னன்னே தெரியல!! 😲</em> | |
| </p> | |
| </div>""") | |
| gr.Markdown("### Step 2 — Choose a model\n" | |
| "Not sure which to pick? **Use SetFit — it is the most accurate.**") | |
| with gr.Row(): | |
| t1_model = gr.Radio(choices=MODEL_CHOICES, value=MODEL_CHOICES[0], | |
| label="Which model should analyse your tweet?") | |
| gr.HTML(""" | |
| <div style="display:flex; gap:12px; flex-wrap:wrap; margin-bottom:16px;"> | |
| <div style="flex:1; min-width:180px; background:#eff6ff; border:1px solid #bfdbfe; | |
| border-radius:10px; padding:12px;"> | |
| <p style="margin:0; font-weight:700; color:#1e40af;">⭐ SetFit (Recommended)</p> | |
| <p style="margin:4px 0 0; font-size:0.85em; color:#1e3a8a;"> | |
| Accuracy: <strong>80.65%</strong><br>Best for everyday use. | |
| </p> | |
| </div> | |
| <div style="flex:1; min-width:180px; background:#f0fdf4; border:1px solid #bbf7d0; | |
| border-radius:10px; padding:12px;"> | |
| <p style="margin:0; font-weight:700; color:#166534;">Prompt-Based</p> | |
| <p style="margin:4px 0 0; font-size:0.85em; color:#14532d;"> | |
| Accuracy: <strong>58.71%</strong><br>No training needed. Handles noisy text well. | |
| </p> | |
| </div> | |
| </div>""") | |
| t1_btn = gr.Button("Analyse Emotion", variant="primary", size="lg") | |
| gr.Markdown("### Results") | |
| t1_result = gr.Markdown(value="_Your results will appear here after you click Analyse._") | |
| with gr.Row(): | |
| t1_table = gr.Dataframe(label="Full breakdown — all six emotions and their confidence scores", wrap=True) | |
| t1_chart = gr.Image(label="Confidence chart — how sure is the model about each emotion?", height=320) | |
| t1_btn.click(fn=tab_single, inputs=[t1_text, t1_model], | |
| outputs=[t1_result, t1_table, t1_chart]) | |
| # ── TAB 2 ───────────────────────────────────────────── | |
| with gr.Tab("Compare All Models"): | |
| gr.Markdown("### See what each model thinks about the same tweet\n" | |
| "Runs your tweet through both models at once for a side-by-side comparison.") | |
| t2_text = gr.Textbox(label="Your tweet", lines=3, | |
| placeholder="Type any Sinhala-Tamil tweet here...") | |
| t2_btn = gr.Button("Compare All Models", variant="primary", size="lg") | |
| t2_result = gr.Markdown(value="_Results will appear here after you click Compare._") | |
| with gr.Row(): | |
| t2_table = gr.Dataframe(label="Side-by-side comparison", wrap=True) | |
| t2_chart = gr.Image(label="Visual comparison — colour shows the predicted emotion", height=320) | |
| t2_btn.click(fn=tab_compare, inputs=[t2_text], | |
| outputs=[t2_result, t2_table, t2_chart]) | |
| # TAB 3 | |
| with gr.Tab("Analyse Many Tweets"): | |
| gr.Markdown("### Analyse a whole spreadsheet of tweets at once\n" | |
| "Upload a file and STEMO will classify each tweet automatically.") | |
| gr.HTML(""" | |
| <div style="background:#eff6ff; border:1px solid #bfdbfe; border-radius:10px; | |
| padding:14px 16px; margin-bottom:16px;"> | |
| <p style="margin:0; font-size:0.9em; color:#1e40af;"> | |
| <strong>📋 How to prepare your file:</strong><br> | |
| • <strong>.xlsx (Excel)</strong> or <strong>.csv</strong> format<br> | |
| • Must have a column containing the tweets (default name: <strong>text</strong>)<br> | |
| • Other columns are kept in the results | |
| </p> | |
| </div>""") | |
| with gr.Row(): | |
| with gr.Column(): | |
| t3_file = gr.File(label="Upload your file (.xlsx or .csv)", | |
| file_types=['.xlsx', '.csv']) | |
| t3_text_col = gr.Textbox(label="Tweet column name", value="text", | |
| info="Column in your file that contains the tweets") | |
| t3_model = gr.Dropdown(choices=MODEL_CHOICES, value=MODEL_CHOICES[0], | |
| label="Which model to use?", | |
| info="SetFit is recommended for best accuracy") | |
| t3_btn = gr.Button("Start Analysis", variant="primary", size="lg") | |
| t3_result = gr.Markdown(value="_Upload a file and click Start Analysis to begin._") | |
| t3_table = gr.Dataframe(label="Preview — first 10 rows of results", wrap=True) | |
| t3_download = gr.File(label="Download Full Results (.xlsx)", visible=False) | |
| def run_batch(file, model_choice, text_col): | |
| status, preview, path = tab_batch(file, model_choice, text_col) | |
| return status, preview, gr.update(value=path, visible=path is not None) | |
| t3_btn.click(fn=run_batch, inputs=[t3_file, t3_model, t3_text_col], | |
| outputs=[t3_result, t3_table, t3_download]) | |
| # ── TAB 4 ───────────────────────────────────────────── | |
| with gr.Tab("Evaluate Model Performance"): | |
| gr.Markdown("### For researchers — test how well the model performs on your labelled data\n" | |
| "Upload a file with tweets and correct emotion labels for a full performance report.") | |
| gr.HTML(""" | |
| <div style="background:#fdf4ff; border:1px solid #e9d5ff; border-radius:10px; | |
| padding:14px 16px; margin-bottom:16px;"> | |
| <p style="margin:0; font-size:0.9em; color:#6b21a8;"> | |
| <strong>Your file needs two columns:</strong><br> | |
| • <strong>text</strong> — the tweet<br> | |
| • <strong>label</strong> — the correct emotion | |
| (Happy, Anger, Sadness, Fear, Surprise, Disgust) | |
| </p> | |
| </div>""") | |
| with gr.Row(): | |
| with gr.Column(): | |
| t4_file = gr.File(label="Upload labelled test file (.xlsx or .csv)", | |
| file_types=['.xlsx', '.csv']) | |
| with gr.Row(): | |
| t4_text_col = gr.Textbox(label="Tweet column name", value="text") | |
| t4_label_col = gr.Textbox(label="Label column name", value="label") | |
| t4_model = gr.Dropdown(choices=MODEL_CHOICES, value=MODEL_CHOICES[0], | |
| label="Which model to evaluate?") | |
| t4_btn = gr.Button("Run Evaluation", variant="primary", size="lg") | |
| t4_result = gr.Markdown(value="_Upload a labelled file and click Run Evaluation._") | |
| with gr.Row(): | |
| t4_table = gr.Dataframe(label="Performance metrics", wrap=True) | |
| t4_chart = gr.Image(label="Confusion matrix", height=400) | |
| t4_download = gr.File( | |
| label="Download Full Report (.xlsx) — includes predictions and per-emotion F1", | |
| visible=False) | |
| def run_eval(file, model_choice, text_col, label_col): | |
| report, metrics, img, path = tab_evaluate(file, model_choice, text_col, label_col) | |
| return report, metrics, img, gr.update(value=path, visible=path is not None) | |
| t4_btn.click(fn=run_eval, | |
| inputs=[t4_file, t4_model, t4_text_col, t4_label_col], | |
| outputs=[t4_result, t4_table, t4_chart, t4_download]) | |
| gr.HTML(""" | |
| <div style="text-align:center; margin-top:32px; padding:20px; | |
| background:#f8fafc; border-radius:12px; border:1px solid #e2e8f0;"> | |
| <p style="margin:0; color:#475569; font-size:0.95em;"> | |
| <strong>STEMO</strong> — Sinhala-Tamil Emotion Model | |
| </p> | |
| <p style="margin:6px 0 0; color:#94a3b8; font-size:0.85em;"> | |
| SetFit (80.65%) | Prompt-Based XLM-RoBERTa (58.71%)<br> | |
| Trained on 1,013 code-mixed Sinhala-Tamil tweets from the ACTSEA corpus. | |
| </p> | |
| </div> | |
| """) | |
| demo.launch() |