Update app.py

app.py

@@ -1,112 +1,86 @@
 import gradio as gr
 from transformers import pipeline
 
-# Initialize the classifier
-classifier = pipeline("zero-shot-classification", model="tasksource/ModernBERT-base-nli")
+# 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")
 
-def zeroShotClassification(text_input, candidate_labels):
-    # Clean and process the labels
-    labels = [label.strip() for label in candidate_labels.split(',')]
+def process_input(text_input, labels_or_premise, mode):
+    if mode == "Zero-Shot Classification":
+        # Clean and process the labels
+        labels = [label.strip() for label in labels_or_premise.split(',')]
+        
+        # Get predictions
+        prediction = zero_shot_classifier(text_input, labels)
+        results = {label: score for label, score in zip(prediction['labels'], prediction['scores'])}
+        return results, ''
     
-    # Get predictions
-    prediction = classifier(text_input, labels)
+    else:  # NLI mode
+        # Process as premise-hypothesis pair
+        prediction = nli_classifier([{"text": text_input, "text_pair": labels_or_premise}])
+        results = {pred['label']: pred['score'] for pred in prediction}
+        return results, ''
+
+# Create the interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 ModernBERT Text Analysis")
     
-    # For Label component: Return raw scores (not percentage strings)
-    results = {label: score for label, score in zip(prediction['labels'], prediction['scores'])}
+    mode = gr.Radio(
+        ["Zero-Shot Classification", "Natural Language Inference"],
+        label="Select Mode",
+        value="Zero-Shot Classification"
+    )
     
+    with gr.Column():
+        text_input = gr.Textbox(
+            label="✍️ Input Text",
+            placeholder="Enter your text...",
+            lines=3
+        )
+        
+        labels_or_premise = gr.Textbox(
+            label="🏷️ Categories / Premise",
+            placeholder="Enter comma-separated categories or premise text...",
+            lines=2
+        )
+        
+        submit_btn = gr.Button("Submit")
+        
+        outputs = [
+            gr.Label(label="📊 Results"),
+            gr.Markdown(label="📈 Analysis", visible=False)
+        ]
 
-    return results, ''
+    # Different examples for each mode
+    zero_shot_examples = [
+        ["I need to buy groceries", "shopping, urgent tasks, leisure, philosophy"],
+        ["The sun is very bright today", "weather, astronomy, complaints, poetry"],
+        ["I love playing video games", "entertainment, sports, education, business"],
+        ["The car won't start", "transportation, art, cooking, literature"],
+        ["She wrote a beautiful poem", "creativity, finance, exercise, technology"]
+    ]
 
-# Create the interface
-demo = gr.Interface(
-    fn=zeroShotClassification,
-    inputs=[
-        gr.Textbox(
-            label="✍️ Input Text",
-            placeholder="Enter the text you want to classify...",
-            lines=3,
-            elem_classes=["example-text"]
-        ),
-        gr.Textbox(
-            label="🏷️ Category Labels",
-            placeholder="Enter comma-separated categories (e.g., happy, sad, excited, confused)",
-            lines=2,
-            elem_classes=["example-text"]
+    nli_examples = [
+        ["A man is sleeping on a couch", "The man is awake"],
+        ["The restaurant is full of people", "The place is empty"],
+        ["The child is playing with toys", "The kid is having fun"],
+        ["It's raining outside", "The weather is wet"],
+        ["The dog is barking at the mailman", "There is a cat"]
+    ]
+
+    def update_examples(mode_value):
+        return gr.Examples(
+            zero_shot_examples if mode_value == "Zero-Shot Classification" else nli_examples,
+            inputs=[text_input, labels_or_premise]
         )
-    ],
-    outputs=[
-        gr.Label(label="📊 Classification Results"),
-        gr.Markdown(label="📈 Detailed Analysis", elem_classes=["markdown-text"])
-    ],
-    title="🤖 Zero-Shot Text Classification with ModernBERT",
-    description="""
-    Classify any text into arbirary categories or perform natural language inference with ModernBERT 
-    
-    **How to use:**
-    1. Enter your text in the first box
-    2. Add comma-separated category labels in the second box
-    3. Click submit to see how your text matches each category
+
+    mode.change(fn=update_examples, inputs=[mode], outputs=gr.Examples())
     
-    Try the examples below or create your own classifications!
-    """,
-    examples=[
-        ["One day I will see the world", "travel, adventure, dreams, future"],
-        ["The movie had amazing special effects but a weak plot", "entertainment, technology, criticism, story"],
-        ["This new phone has an amazing camera and great battery life", "technology, photography, consumer, review"],
-        ["Mix flour, sugar, and eggs until well combined", "cooking, baking, instructions, food"],
-        ["Scientists discovered a new species of butterfly in the Amazon", "science, nature, discovery, environment"],
-        ["The team scored in the final minute to win the championship", "sports, victory, competition, excitement"],
-        ["The painting uses vibrant colors to express deep emotions", "art, emotion, creativity, analysis"]
-    ],
-    cache_examples=False,
-    css="""
-    footer {display:none !important}
-    .output-markdown{display:none !important}
-    .gradio-container {
-        font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
-        max-width: 1200px !important;
-    }
-    .gr-button-primary {
-        background: linear-gradient(90deg, #11142D, #253885) !important;
-        border: none !important;
-        color: white !important;
-        border-radius: 12px !important;
-        transition: all 0.3s ease !important;
-    }
-    .gr-button-primary:hover {
-        transform: translateY(-2px) !important;
-        box-shadow: 0 4px 12px rgba(17, 20, 45, 0.3) !important;
-        background: linear-gradient(90deg, #253885, #4285F4) !important;
-    }
-    .gr-input, .gr-textarea {
-        border-radius: 8px !important;
-        border: 2px solid #E2E8F0 !important;
-        padding: 12px !important;
-        font-size: 16px !important;
-    }
-    .gr-input:focus, .gr-textarea:focus {
-        border-color: #253885 !important;
-        box-shadow: 0 0 0 3px rgba(37, 56, 133, 0.2) !important;
-    }
-    .gr-panel {
-        border-radius: 16px !important;
-        box-shadow: 0 4px 15px -1px rgba(0, 0, 0, 0.1) !important;
-        background: white !important;
-    }
-    .gr-box {
-        border-radius: 12px !important;
-        background: white !important;
-    }
-    .markdown-text {
-        font-size: 16px !important;
-        line-height: 1.6 !important;
-    }
-    .example-text {
-        font-family: 'Inter', sans-serif !important;
-        color: #11142D !important;
-    }
-    """
-)
+    submit_btn.click(
+        fn=process_input,
+        inputs=[text_input, labels_or_premise, mode],
+        outputs=outputs
+    )
 
 # Launch the app
 if __name__ == "__main__":