Update app.py

app.py

@@ -44,40 +44,60 @@ except ImportError:
 # Configuration
 # -----------------------------------------------------------------------------
 MODEL_NAME = "msmaje/phdhatamodel"
-SUPPORTED_LANGUAGES = ["Hausa", "Yoruba", "Igbo", "Swahili", "Amharic", "Nigerian Pidgin"]
+SUPPORTED_LANGUAGES = ["Hausa", "Yoruba", "Igbo", "Nigerian Pidgin"]
 LANGUAGE_CODES = {
     "Hausa": "ha",
     "Yoruba": "yo", 
     "Igbo": "ig",
-    "Swahili": "sw",
-    "Amharic": "am",
     "Nigerian Pidgin": "pcm"
 }
 
 # -----------------------------------------------------------------------------
 # Model Loading
 # -----------------------------------------------------------------------------
-print("Loading model and tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-model.eval()
-print("✅ Model loaded successfully!")
+print("📥 Loading model and tokenizer...")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        MODEL_NAME,
+        output_attentions=True  # Enable attention outputs for explainability
+    )
+    model.eval()
+    print("✅ Model loaded successfully!")
+    print(f"   Model: {MODEL_NAME}")
+    print(f"   Device: {'GPU' if torch.cuda.is_available() else 'CPU'}")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    raise
 
 # Initialize explainability tools
 if LIME_AVAILABLE:
-    lime_explainer = LimeTextExplainer(class_names=["Human", "AI"])
+    try:
+        lime_explainer = LimeTextExplainer(class_names=["Human", "AI"])
+        print("✅ LIME explainer initialized")
+    except Exception as e:
+        print(f"⚠️  LIME initialization failed: {e}")
+        LIME_AVAILABLE = False
 
 if SHAP_AVAILABLE:
-    # Create a wrapper for SHAP
-    def model_predict_proba(texts):
-        inputs = tokenizer(texts, return_tensors="pt", truncation=True, 
-                          max_length=128, padding=True)
-        with torch.no_grad():
-            outputs = model(**inputs)
-            probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        return probs.numpy()
-    
-    shap_explainer = shap.Explainer(model_predict_proba, tokenizer)
+    try:
+        # Create a wrapper for SHAP
+        def model_predict_proba(texts):
+            if isinstance(texts, str):
+                texts = [texts]
+            inputs = tokenizer(texts, return_tensors="pt", truncation=True, 
+                              max_length=128, padding=True)
+            with torch.no_grad():
+                outputs = model(**inputs)
+                probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+            return probs.numpy()
+        
+        shap_explainer = shap.Explainer(model_predict_proba, tokenizer)
+        print("✅ SHAP explainer initialized")
+    except Exception as e:
+        print(f"⚠️  SHAP initialization failed: {e}")
+        print("   Will use attention-based explanations as fallback")
+        SHAP_AVAILABLE = False
 
 # -----------------------------------------------------------------------------
 # Bias and Fairness Metrics
@@ -153,39 +173,50 @@ def get_shap_explanation(text, language="English"):
         return "⚠️ SHAP is not installed. Install with: pip install shap", None
     
     try:
-        # Get SHAP values
-        shap_values = shap_explainer([text])
+        # Simpler approach - use attention weights as proxy for SHAP
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
         
-        # Create visualization
-        fig, ax = plt.subplots(figsize=(12, 6))
-        shap.plots.text(shap_values[0], display=False)
-        plt.tight_layout()
+        with torch.no_grad():
+            outputs = model(**inputs, output_attentions=True)
+            # Get mean attention across all layers and heads
+            attentions = outputs.attentions
+            mean_attention = torch.mean(torch.stack([att.mean(dim=1) for att in attentions]), dim=0)
+            token_importance = mean_attention[0].sum(dim=0).numpy()
         
-        # Extract token attributions
-        tokens = tokenizer.tokenize(text)[:20]  # Limit to first 20 tokens
-        values = shap_values.values[0][:len(tokens), 1]  # AI class
+        # Get tokens
+        tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
+        tokens = tokens[1:-1]  # Remove [CLS] and [SEP]
+        token_importance = token_importance[1:-1]  # Match tokens
         
-        attribution_data = {
-            "Token": tokens,
-            "Attribution": values.tolist()
-        }
+        # Normalize
+        token_importance = token_importance / (token_importance.max() + 1e-8)
         
-        explanation = f"## SHAP Explanation for {language}\n\n"
-        explanation += "Tokens with **positive values** push toward AI-generated classification.\n"
-        explanation += "Tokens with **negative values** push toward Human-written classification.\n\n"
-        explanation += f"Top 5 most influential tokens:\n"
+        # Create simple bar plot
+        fig, ax = plt.subplots(figsize=(12, 6))
+        colors = ['red' if x < 0 else 'green' for x in token_importance]
+        ax.barh(range(min(20, len(tokens))), token_importance[:20], color=colors[:20])
+        ax.set_yticks(range(min(20, len(tokens))))
+        ax.set_yticklabels(tokens[:20])
+        ax.set_xlabel('Importance (Attention Weight)')
+        ax.set_title(f'Token Importance - {language}')
+        ax.invert_yaxis()
+        plt.tight_layout()
+        
+        explanation = f"## Attention-Based Explanation for {language}\n\n"
+        explanation += "Tokens with **higher values** are more important for classification.\n\n"
+        explanation += f"Top 5 most important tokens:\n"
         
-        top_indices = np.argsort(np.abs(values))[-5:][::-1]
+        top_indices = np.argsort(token_importance)[-5:][::-1]
         for idx in top_indices:
-            token = tokens[idx]
-            value = values[idx]
-            direction = "→ AI" if value > 0 else "→ Human"
-            explanation += f"- **{token}**: {value:.4f} {direction}\n"
+            if idx < len(tokens):
+                token = tokens[idx]
+                value = token_importance[idx]
+                explanation += f"- **{token}**: {value:.4f}\n"
         
-        return explanation, (fig, attribution_data)
+        return explanation, fig
     
     except Exception as e:
-        return f"❌ SHAP explanation failed: {str(e)}", None
+        return f"❌ Explanation failed: {str(e)}", None
 
 def get_lime_explanation(text, language="English"):
     """Generate LIME-based explanation"""
@@ -194,19 +225,27 @@ def get_lime_explanation(text, language="English"):
     
     try:
         def predict_fn(texts):
-            inputs = tokenizer(texts, return_tensors="pt", truncation=True,
-                             max_length=128, padding=True)
-            with torch.no_grad():
-                outputs = model(**inputs)
-                probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-            return probs.numpy()
+            """Prediction function for LIME"""
+            if isinstance(texts, str):
+                texts = [texts]
+            
+            results = []
+            for txt in texts:
+                inputs = tokenizer(txt, return_tensors="pt", truncation=True,
+                                 max_length=128, padding=True)
+                with torch.no_grad():
+                    outputs = model(**inputs)
+                    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+                results.append(probs[0].numpy())
+            
+            return np.array(results)
         
         # Generate explanation
         exp = lime_explainer.explain_instance(
             text, 
             predict_fn, 
             num_features=10,
-            num_samples=100
+            num_samples=50  # Reduced for speed
         )
         
         # Create visualization
@@ -219,7 +258,7 @@ def get_lime_explanation(text, language="English"):
         explanation = f"## LIME Explanation for {language}\n\n"
         explanation += "Features with **positive weights** indicate AI-generated characteristics.\n"
         explanation += "Features with **negative weights** indicate Human-written characteristics.\n\n"
-        explanation += "Top contributing features:\n"
+        explanation += "Top contributing features:\n\n"
         
         for feature, weight in weights[:5]:
             direction = "→ AI" if weight > 0 else "→ Human"
@@ -228,7 +267,7 @@ def get_lime_explanation(text, language="English"):
         return explanation, fig
     
     except Exception as e:
-        return f"❌ LIME explanation failed: {str(e)}", None
+        return f"❌ LIME explanation failed: {str(e)}\n\nTry using SHAP instead.", None
 
 # -----------------------------------------------------------------------------
 # Main Classification Function
@@ -263,29 +302,37 @@ def classify_with_explanation(text, language, explainer_type="SHAP"):
     else:
         result += "❓ **Low confidence** - Uncertain, mixed characteristics detected\n"
     
-    # Probability breakdown
-    prob_chart = {
+    # Probability breakdown - Create DataFrame for BarPlot
+    prob_data = pd.DataFrame({
         "Class": ["Human-written", "AI-generated"],
         "Probability": [float(probabilities[0][0]), float(probabilities[0][1])]
-    }
+    })
     
     # Generate explanation
-    explanation_text = None
+    explanation_text = ""
     explanation_viz = None
     
     if explainer_type == "SHAP" and SHAP_AVAILABLE:
         explanation_text, explanation_viz = get_shap_explanation(text, language)
+        if explanation_viz and isinstance(explanation_viz, tuple):
+            explanation_viz = explanation_viz[0]  # Extract just the figure
     elif explainer_type == "LIME" and LIME_AVAILABLE:
         explanation_text, explanation_viz = get_lime_explanation(text, language)
     elif explainer_type == "Both":
         shap_text, shap_viz = get_shap_explanation(text, language)
         lime_text, lime_viz = get_lime_explanation(text, language)
         explanation_text = shap_text + "\n\n---\n\n" + lime_text
-        explanation_viz = (shap_viz, lime_viz) if shap_viz and lime_viz else shap_viz or lime_viz
+        # Use SHAP visualization by default for "Both"
+        if shap_viz and isinstance(shap_viz, tuple):
+            explanation_viz = shap_viz[0]
+        elif isinstance(shap_viz, plt.Figure):
+            explanation_viz = shap_viz
+        else:
+            explanation_viz = lime_viz
     else:
-        explanation_text = "⚠️ Selected explainer not available"
+        explanation_text = "⚠️ Selected explainer not available. Please install SHAP and/or LIME."
     
-    return result, prob_chart, explanation_text, explanation_viz
+    return result, prob_data, explanation_text, explanation_viz
 
 # -----------------------------------------------------------------------------
 # Bias Auditing Function
@@ -431,12 +478,28 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                         x="Class",
                         y="Probability",
                         title="Prediction Probabilities",
-                        y_lim=[0, 1]
+                        y_lim=[0, 1],
+                        height=300,
+                        width=400
                     )
             
             with gr.Row():
-                explanation_output = gr.Markdown(label="Explanation")
-                explanation_viz = gr.Plot(label="Visual Explanation")
+                with gr.Column():
+                    explanation_output = gr.Markdown(label="Explanation")
+                with gr.Column():
+                    explanation_viz = gr.Plot(label="Visual Explanation")
+            
+            # Examples to help users
+            gr.Examples(
+                examples=[
+                    ["Ka rubuta labari game da kasuwa a Kano", "Hausa", "SHAP"],
+                    ["Ìwé yìí jẹ́ ìwé tó dára púpọ̀ fún àwọn akẹ́kọ̀ọ́", "Yoruba", "LIME"],
+                    ["Akwụkwọ a dị mma maka ụmụ akwụkwọ", "Igbo", "SHAP"],
+                    ["Dis book dey very good for students wey wan learn", "Nigerian Pidgin", "Both"]
+                ],
+                inputs=[text_input, language_select, explainer_select],
+                label="Try these examples in different languages"
+            )
             
             classify_btn.click(
                 fn=classify_with_explanation,
@@ -491,7 +554,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
             - Per-language performance metrics
             
             ## 🌍 Supported Languages
-            Hausa, Yoruba, Igbo, Swahili, Amharic, Nigerian Pidgin
+            Hausa, Yoruba, Igbo, Nigerian Pidgin
             
             ## 📊 Model Performance
             - Accuracy: 100%
@@ -500,9 +563,11 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
             - AAOD: 0.0 (No bias)
             
             ## 🔬 Technical Details
-            - Base Model: AfroXLMR-base
+            - Base Model: AfroXLMR-base (davlan/afro-xlmr-base)
             - Parameters: ~270M
             - Max Sequence Length: 128 tokens
+            - Training Dataset: PhD HATA African Dataset
+            - Languages: 4 West African languages
             
             ## 📚 Citation
             ```bibtex
@@ -524,11 +589,4 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
     """)
 
 if __name__ == "__main__":
-    demo.queue(api_open=False)
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        share=True   # <-- important for Spaces
-    )
-
+    demo.launch()