Update app.py

app.py

@@ -2,8 +2,8 @@ import gradio as gr
 from transformers import pipeline
 import re
 
+# Custom sentence tokenizer
 def sent_tokenize(text):
-    # Regular expression to split sentences
     sentence_endings = re.compile(r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?|!)(\s|$)')
     sentences = sentence_endings.split(text)
     return [s.strip() for s in sentences if s.strip()]
@@ -34,6 +34,70 @@ long_context_examples = [
      "The cafe is experiencing a slow, quiet morning"]
 ]
 
+def get_label_color(label):
+    """Return color based on NLI label."""
+    colors = {
+        'ENTAILMENT': '#90EE90',  # Light green
+        'NEUTRAL': '#FFE5B4',     # Peach
+        'CONTRADICTION': '#FFB6C1' # Light pink
+    }
+    return colors.get(label, '#FFFFFF')
+
+def create_analysis_html(sentence_results, global_label):
+    """Create HTML table for sentence analysis with color coding."""
+    html = """
+    <style>
+        .analysis-table {
+            width: 100%;
+            border-collapse: collapse;
+            margin: 20px 0;
+            font-family: Arial, sans-serif;
+        }
+        .analysis-table th, .analysis-table td {
+            padding: 12px;
+            border: 1px solid #ddd;
+            text-align: left;
+        }
+        .analysis-table th {
+            background-color: #f5f5f5;
+        }
+        .global-prediction {
+            padding: 15px;
+            margin: 20px 0;
+            border-radius: 5px;
+            font-weight: bold;
+        }
+    </style>
+    """
+    
+    # Add global prediction box
+    html += f"""
+    <div class="global-prediction" style="background-color: {get_label_color(global_label)}">
+        Global Prediction: {global_label}
+    </div>
+    """
+    
+    # Create table
+    html += """
+    <table class="analysis-table">
+        <tr>
+            <th>Sentence</th>
+            <th>Prediction</th>
+        </tr>
+    """
+    
+    # Add rows for each sentence
+    for result in sentence_results:
+        html += f"""
+        <tr style="background-color: {get_label_color(result['prediction'])}">
+            <td>{result['sentence']}</td>
+            <td>{result['prediction']}</td>
+        </tr>
+        """
+    
+    html += "</table>"
+    return html
+
 def process_input(text_input, labels_or_premise, mode):
     if mode == "Zero-Shot Classification":
         labels = [label.strip() for label in labels_or_premise.split(',')]
@@ -48,6 +112,7 @@ def process_input(text_input, labels_or_premise, mode):
         # Global prediction
         global_pred = nli_classifier([{"text": text_input, "text_pair": labels_or_premise}], return_all_scores=True)[0]
         global_results = {pred['label']: pred['score'] for pred in global_pred}
+        global_label = max(global_results.items(), key=lambda x: x[1])[0]
         
         # Sentence-level analysis
         sentences = sent_tokenize(text_input)
@@ -63,50 +128,10 @@ def process_input(text_input, labels_or_premise, mode):
                 'scores': sent_scores
             })
         
-        # Create markdown analysis
-        analysis_md = "## Global Prediction\n"
-        max_global_label = max(global_results.items(), key=lambda x: x[1])[0]
-        analysis_md += f"Overall prediction: **{max_global_label}**\n\n"
-        analysis_md += "## Sentence-Level Analysis\n"
-        
-        for i, result in enumerate(sentence_results, 1):
-            analysis_md += f"\n### Sentence {i}\n"
-            analysis_md += f"*{result['sentence']}*\n"
-            analysis_md += f"Prediction: **{result['prediction']}**\n"
-            scores_str = ", ".join([f"{label}: {score:.2f}" for label, score in result['scores'].items()])
-            analysis_md += f"Scores: {scores_str}\n"
-        
-        return global_results, analysis_md
-
-def update_interface(mode):
-    if mode == "Zero-Shot Classification":
-        return (
-            gr.update(
-                label="🏷️ Categories", 
-                placeholder="Enter comma-separated categories...",
-                value=zero_shot_examples[0][1]
-            ),
-            gr.update(value=zero_shot_examples[0][0])
-        )
-    elif mode == "Natural Language Inference":
-        return (
-            gr.update(
-                label="🔎 Hypothesis", 
-                placeholder="Enter a hypothesis to compare with the premise...",
-                value=nli_examples[0][1]
-            ),
-            gr.update(value=nli_examples[0][0])
-        )
-    else:  # Long Context NLI
-        return (
-            gr.update(
-                label="🔎 Global Hypothesis",
-                placeholder="Enter a hypothesis to test against the full context...",
-                value=long_context_examples[0][1]
-            ),
-            gr.update(value=long_context_examples[0][0])
-        )
+        analysis_html = create_analysis_html(sentence_results, global_label)
+        return global_results, analysis_html
 
+# [Previous interface code remains the same until the outputs definition]
 with gr.Blocks() as demo:
     gr.Markdown("""
     # tasksource/ModernBERT-nli demonstration
@@ -142,7 +167,7 @@ with gr.Blocks() as demo:
         
         outputs = [
             gr.Label(label="📊 Results"),
-            gr.Markdown(label="📈 Sentence Analysis", visible=True)
+            gr.HTML(label="📈 Sentence Analysis")  # Changed from Markdown to HTML
         ]
 
         with gr.Column(variant="panel") as zero_shot_examples_panel: