Update app.py

app.py

@@ -1,35 +1,36 @@
 # app.py
-# Amharic Poetry Joint Theme Classifier - Multi-Model Comparison
+# Amharic Poetry Joint Theme Classifier
+# Based on: Transformer-Based Joint Thematic Classification for Amharic Poetry
 # Compatible with Gradio 6.20.0
 
 import gradio as gr
 import matplotlib.pyplot as plt
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 import os
 import json
 import numpy as np
-from transformers import AutoTokenizer, AutoModel, RobertaConfig
+from transformers import AutoTokenizer, AutoModel
 
 print("=" * 80)
 print("🚀 Amharic Poetry Joint Theme Classifier")
 print("=" * 80)
 print("📊 Model: Rasyosef-RoBERTa (Multi-Task)")
+print("📄 Paper: Transformer-Based Joint Thematic Classification for Amharic Poetry")
 print("=" * 80)
 
 # ============================================
 # CLASS NAMES (English and Amharic)
-# MUST MATCH THE TRAINING ORDER
+# MATCHES THE PAPER'S 6 THEMATIC CATEGORIES
 # ============================================
 
 CLASS_NAMES_EN = [
-    'Religious',      # LABEL_0
-    'Ethical',        # LABEL_1
-    'Political',      # LABEL_2
-    'Philosophical',  # LABEL_3
-    'Historical',     # LABEL_4
-    'Love'            # LABEL_5
+    'Religious',      # LABEL_0 - ሃይማኖታዊ ግጥም
+    'Ethical',        # LABEL_1 - ሥነ-ምግባራዊ ግጥም
+    'Political',      # LABEL_2 - ፖለቲካዊ ግጥም
+    'Philosophical',  # LABEL_3 - ፍልስፍናዊ ግጥም
+    'Historical',     # LABEL_4 - ታሪካዊ ግጥም
+    'Love'            # LABEL_5 - የፍቅር ግጥም
 ]
 
 CLASS_NAMES_AM = [
@@ -47,29 +48,35 @@ LABEL_TO_ID = {name: i for i, name in enumerate(CLASS_NAMES_EN)}
 
 # ============================================
 # MULTI-TASK MODEL ARCHITECTURE
-# MATCHES THE TRAINED MODEL
+# MATCHES THE PAPER'S JOINT LEARNING FRAMEWORK
 # ============================================
 
 class MultiTaskAmharicPoetryModel(nn.Module):
-    def __init__(self, model_name, num_labels, dropout=0.3, projection_dim=128):
+    """Multi-task model for joint dominant-theme and multi-label classification.
+    This architecture matches the paper's joint learning framework."""
+    
+    def __init__(self, model_name, num_labels=6, dropout=0.3, projection_dim=128):
         super().__init__()
         self.num_labels = num_labels
         self.projection_dim = projection_dim
 
-        # Shared encoder
+        # Shared encoder (Rasyosef-RoBERTa)
         self.encoder = AutoModel.from_pretrained(model_name)
         hidden_size = self.encoder.config.hidden_size
 
-        # Dropout
+        # Dropout for regularization
         self.dropout = nn.Dropout(dropout)
 
         # Multi-label head (Sigmoid, 6 outputs)
+        # Used for multi-label thematic classification
         self.multi_label_head = nn.Linear(hidden_size, num_labels)
 
         # Main-class head (Softmax, 6 classes)
+        # Used for dominant-theme prediction
         self.main_class_head = nn.Linear(hidden_size, num_labels)
 
         # Contrastive learning projection head
+        # Used for discriminative feature learning
         self.projection_head = nn.Sequential(
             nn.Linear(hidden_size, hidden_size),
             nn.ReLU(),
@@ -83,13 +90,13 @@ class MultiTaskAmharicPoetryModel(nn.Module):
             input_ids=input_ids,
             attention_mask=attention_mask
         )
-        # Use CLS token representation
+        # Use CLS token representation for document-level semantics
         pooled_output = outputs.last_hidden_state[:, 0, :]
         pooled_output = self.dropout(pooled_output)
 
         # Task-specific outputs
-        multi_logits = self.multi_label_head(pooled_output)
-        main_logits = self.main_class_head(pooled_output)
+        multi_logits = self.multi_label_head(pooled_output)  # Sigmoid for multi-label
+        main_logits = self.main_class_head(pooled_output)    # Softmax for dominant theme
 
         if return_projection:
             projection = self.projection_head(pooled_output)
@@ -112,28 +119,25 @@ def load_model():
     if model is not None:
         return model, tokenizer, device
 
-    device = torch.device(
-        "cuda" if torch.cuda.is_available() else "cpu"
-    )
-
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"📱 Device: {device}")
 
     if not os.path.exists("best_model.bin"):
         print("❌ best_model.bin not found!")
-        print(os.listdir("."))
+        print("📁 Files:", os.listdir("."))
         return None, None, device
 
     try:
-        print("📥 Loading tokenizer from current directory...")
+        print("📥 Loading tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained(".")
-
+        
         # Ensure pad token exists
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
         print("🏗️ Building Multi-Task Model architecture...")
 
-        # Build the multi-task model
+        # Build the multi-task model with Rasyosef-RoBERTa
         model = MultiTaskAmharicPoetryModel(
             "rasyosef/roberta-base-amharic",
             num_labels=6,
@@ -143,33 +147,26 @@ def load_model():
 
         print("📦 Loading best_model.bin weights...")
 
-        state_dict = torch.load(
-            "best_model.bin",
-            map_location=device
-        )
+        state_dict = torch.load("best_model.bin", map_location=device)
 
-        # Load state dict
-        missing, unexpected = model.load_state_dict(
-            state_dict,
-            strict=False
-        )
+        # Load state dict with strict=False to handle mismatches
+        missing, unexpected = model.load_state_dict(state_dict, strict=False)
 
         if missing:
-            print(f"⚠️ Missing keys: {len(missing)} keys")
+            print(f"⚠️ Missing keys: {len(missing)}")
             if len(missing) <= 10:
-                print("Missing keys:", missing)
+                print("Missing:", missing)
         if unexpected:
-            print(f"⚠️ Unexpected keys: {len(unexpected)} keys")
+            print(f"⚠️ Unexpected keys: {len(unexpected)}")
             if len(unexpected) <= 10:
-                print("Unexpected keys:", unexpected)
+                print("Unexpected:", unexpected)
 
         model.to(device)
         model.eval()
 
-        print("✅ best_model.bin loaded successfully!")
-        print(
-            f"📊 Parameters: {sum(p.numel() for p in model.parameters()):,}"
-        )
+        print(f"✅ Model loaded successfully!")
+        print(f"📊 Parameters: {sum(p.numel() for p in model.parameters()):,}")
+        print(f"📊 Hidden size: {model.encoder.config.hidden_size}")
 
         return model, tokenizer, device
 
@@ -188,15 +185,14 @@ try:
         print("❌ Model loading failed!")
 except Exception as e:
     print(f"❌ Model loading failed: {e}")
-    import traceback
-    traceback.print_exc()
 
 # ============================================
 # PREDICTION FUNCTION
 # ============================================
 
 def gradio_predict(text, threshold=0.5):
-    """Gradio prediction function."""
+    """Gradio prediction function.
+    Uses sigmoid for multi-label and softmax for dominant theme."""
     
     if not text or not text.strip():
         return "⚠️ Please enter a poem.", None
@@ -205,7 +201,7 @@ def gradio_predict(text, threshold=0.5):
         return "❌ Model not loaded. Please check the logs.", None
     
     try:
-        # Tokenize with max_length=128 (from config)
+        # Tokenize with max_length=128 (from paper's config)
         inputs = tokenizer(
             text,
             return_tensors='pt',
@@ -219,17 +215,22 @@ def gradio_predict(text, threshold=0.5):
         # Get predictions
         with torch.no_grad():
             multi_logits, main_logits = model(input_ids, attention_mask)
-            probs = torch.sigmoid(multi_logits)  # Multi-label uses sigmoid
-            probs_np = probs.cpu().numpy()[0]
+            # Multi-label uses sigmoid (as in paper)
+            multi_probs = torch.sigmoid(multi_logits)
+            # Dominant theme uses softmax (as in paper)
+            main_probs = torch.softmax(main_logits, dim=1)
+            
+            multi_probs_np = multi_probs.cpu().numpy()[0]
+            main_probs_np = main_probs.cpu().numpy()[0]
         
-        # Main class (highest probability from multi-label)
-        main_idx = np.argmax(probs_np)
+        # Dominant theme: highest probability from multi-label
+        main_idx = np.argmax(multi_probs_np)
         main_class = ID_TO_LABEL[main_idx]
-        main_confidence = float(probs_np[main_idx])
+        main_confidence = float(multi_probs_np[main_idx])
         
-        # Multi-label (all classes above threshold, excluding main)
+        # Multi-label: all classes above threshold, excluding main
         multi_labels = []
-        for i, prob in enumerate(probs_np):
+        for i, prob in enumerate(multi_probs_np):
             if prob >= threshold and i != main_idx:
                 multi_labels.append({
                     'label': ID_TO_LABEL[i],
@@ -239,10 +240,10 @@ def gradio_predict(text, threshold=0.5):
         
         # All probabilities
         all_probs = {}
-        for i, prob in enumerate(probs_np):
+        for i, prob in enumerate(multi_probs_np):
             all_probs[ID_TO_LABEL[i]] = float(prob)
         
-        # Build output
+        # Build output (matching UI images)
         output = f"## 🎯 Predicted Theme Class of the Poem\n(የግጥሙ ዋና የጭብጥ ምድብ)\n\n"
         
         # Display as "Love poetry (የፍቅር ግጥም)"
@@ -258,7 +259,7 @@ def gradio_predict(text, threshold=0.5):
         else:
             output += "*No additional themes above threshold.*\n"
         
-        # Create plot
+        # Create plot (matching UI images)
         labels = list(all_probs.keys())
         probs = list(all_probs.values())
         
@@ -274,8 +275,9 @@ def gradio_predict(text, threshold=0.5):
         ax.set_ylim(0, 1)
         ax.set_ylabel("Probability", fontsize=12, fontweight='bold')
         ax.set_title("Class Probabilities", fontsize=14, fontweight='bold')
-        ax.axhline(y=threshold, linestyle="--", alpha=0.7, color='#e74c3c', linewidth=2)
+        ax.axhline(y=threshold, linestyle="--", alpha=0.7, color='#e74c3c', linewidth=2, label=f'Threshold ({threshold:.0%})')
         ax.grid(True, alpha=0.2, axis='y')
+        ax.legend(loc='upper right')
         
         for bar, p in zip(bars, probs):
             if p > 0.01:
@@ -292,15 +294,15 @@ def gradio_predict(text, threshold=0.5):
         return f"❌ Error: {str(e)}", None
 
 # ============================================
-# SAMPLE POEMS
+# SAMPLE POEMS (From UI Images and Paper)
 # ============================================
 
 sample_poems = [
-    # Ethical poem
+    # Ethical poem (from UI image - Figure 9)
     """ሰው ከጎረቤቱ አብሮ ለመኖር፣
 በልቡ ይኑረው ደስታና ፍቅር፡፡""",
     
-    # Historical poem
+    # Historical poem (from UI image - Figure 10)
     """መቅደላ አፋፉ ላይ ጩኋት በረከተ
 የሴቱን አናዉቅም ወንድ አንድ ሰዉ ሞተ
 የሰሜኑን ንጉስ ሲንቁ ሲንቁ
@@ -336,7 +338,7 @@ sample_poems = [
 ]
 
 # ============================================
-# GRADIO UI
+# GRADIO UI (Matches UI Images)
 # ============================================
 
 # Custom CSS
@@ -420,7 +422,7 @@ with gr.Blocks(
         outputs=[text_input, threshold_slider, output_text, output_plot]
     )
     
-    # Footer
+    # Footer with paper citation
     gr.Markdown("""
     ---
     <div style="text-align: center; font-size: 13px; color: #5d6d7e;">