Update app.py

app.py

@@ -1,14 +1,14 @@
 import gradio as gr
 from transformers import pipeline
+import nltk
+nltk.download('punkt')
+from nltk.tokenize import sent_tokenize
 
 # Initialize the classifiers
 zero_shot_classifier = pipeline("zero-shot-classification", model="tasksource/ModernBERT-base-nli")
 nli_classifier = pipeline("text-classification", model="tasksource/ModernBERT-base-nli")
 
-if False:
-    gr.load("models/answerdotai/ModernBERT-base").launch()
-
-# Define examples
+# Define examples (including new long context example)
 zero_shot_examples = [
     ["I absolutely love this product, it's amazing!", "positive, negative, neutral"],
     ["I need to buy groceries", "shopping, urgent tasks, leisure, philosophy"],
@@ -25,17 +25,54 @@ nli_examples = [
     ["A German Shepherd is exhibiting defensive behavior towards someone approaching the property", "The animal making noise is feline"]
 ]
 
+long_context_examples = [
+    ["""The small cafe on the corner has been bustling with activity all morning. The aroma of freshly baked pastries wafts through the air, drawing in passersby. The baristas work efficiently behind the counter, crafting intricate latte art. Several customers are seated at wooden tables, engaged in quiet conversations or working on laptops. Through the large windows, sunshine streams in, creating a warm and inviting atmosphere.""",
+     "The cafe is experiencing a slow, quiet morning"]
+]
+
 def process_input(text_input, labels_or_premise, mode):
     if mode == "Zero-Shot Classification":
         labels = [label.strip() for label in labels_or_premise.split(',')]
         prediction = zero_shot_classifier(text_input, labels)
         results = {label: score for label, score in zip(prediction['labels'], prediction['scores'])}
         return results, ''
-    else:  # NLI mode
-        pred= nli_classifier([{"text": text_input, "text_pair": labels_or_premise}],return_all_scores=True)[0]
-        results= {pred['label']:pred['score'] for pred in pred}
-
+    elif mode == "Natural Language Inference":
+        pred = nli_classifier([{"text": text_input, "text_pair": labels_or_premise}], return_all_scores=True)[0]
+        results = {pred['label']: pred['score'] for pred in pred}
         return results, ''
+    else:  # Long Context NLI
+        # 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}
+        
+        # Sentence-level analysis
+        sentences = sent_tokenize(text_input)
+        sentence_results = []
+        
+        for sentence in sentences:
+            sent_pred = nli_classifier([{"text": sentence, "text_pair": labels_or_premise}], return_all_scores=True)[0]
+            sent_scores = {pred['label']: pred['score'] for pred in sent_pred}
+            max_label = max(sent_scores.items(), key=lambda x: x[1])[0]
+            sentence_results.append({
+                'sentence': sentence,
+                'prediction': max_label,
+                '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":
@@ -47,7 +84,7 @@ def update_interface(mode):
             ),
             gr.update(value=zero_shot_examples[0][0])
         )
-    else:
+    elif mode == "Natural Language Inference":
         return (
             gr.update(
                 label="🔎 Hypothesis", 
@@ -56,19 +93,28 @@ def update_interface(mode):
             ),
             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])
+        )
 
 with gr.Blocks() as demo:
     gr.Markdown("""
     # tasksource/ModernBERT-nli demonstration
     
-    This spaces uses [tasksource/ModernBERT-base-nli](https://huggingface.co/tasksource/ModernBERT-base-nli), 
+    This space uses [tasksource/ModernBERT-base-nli](https://huggingface.co/tasksource/ModernBERT-base-nli), 
     fine-tuned from [answerdotai/ModernBERT-base](https://huggingface.co/answerdotai/ModernBERT-base) 
     on tasksource classification tasks. 
     This NLI model achieves high accuracy on logical reasoning and long-context NLI, outperforming Llama 3 8B on ConTRoL and FOLIO.
     """)
 
     mode = gr.Radio(
-        ["Zero-Shot Classification", "Natural Language Inference"],
+        ["Zero-Shot Classification", "Natural Language Inference", "Long Context NLI"],
         label="Select Mode",
         value="Zero-Shot Classification"
     )
@@ -78,21 +124,21 @@ with gr.Blocks() as demo:
             label="✍️ Input Text",
             placeholder="Enter your text...",
             lines=3,
-            value=zero_shot_examples[0][0]  # Initial value
+            value=zero_shot_examples[0][0]
         )
         
         labels_or_premise = gr.Textbox(
             label="🏷️ Categories",
             placeholder="Enter comma-separated categories...",
             lines=2,
-            value=zero_shot_examples[0][1]  # Initial value
+            value=zero_shot_examples[0][1]
         )
         
         submit_btn = gr.Button("Submit")
         
         outputs = [
             gr.Label(label="📊 Results"),
-            gr.Markdown(label="📈 Analysis", visible=False)
+            gr.Markdown(label="📈 Sentence Analysis", visible=True)
         ]
 
         with gr.Column(variant="panel") as zero_shot_examples_panel:
@@ -107,12 +153,20 @@ with gr.Blocks() as demo:
                 examples=nli_examples,
                 inputs=[text_input, labels_or_premise],
                 label="Natural Language Inference Examples",
-        )
+            )
+            
+        with gr.Column(variant="panel") as long_context_examples_panel:
+            gr.Examples(
+                examples=long_context_examples,
+                inputs=[text_input, labels_or_premise],
+                label="Long Context NLI Examples",
+            )
 
     def update_visibility(mode):
         return (
             gr.update(visible=(mode == "Zero-Shot Classification")),
-            gr.update(visible=(mode == "Natural Language Inference"))
+            gr.update(visible=(mode == "Natural Language Inference")),
+            gr.update(visible=(mode == "Long Context NLI"))
         )
 
     mode.change(
@@ -124,7 +178,7 @@ with gr.Blocks() as demo:
     mode.change(
         fn=update_visibility,
         inputs=[mode],
-        outputs=[zero_shot_examples_panel, nli_examples_panel]
+        outputs=[zero_shot_examples_panel, nli_examples_panel, long_context_examples_panel]
     )
     
     submit_btn.click(