Update application

app.py

@@ -12,21 +12,25 @@ from utils import (
     cross_encoder_predict_relevance
 )
 
-def predict(system_prompt, user_prompt, selected_model):
-    if selected_model == "jinaai/jina-embeddings-v2-small-en":
-        predicted_label, probabilities = embeddings_predict_relevance(system_prompt, user_prompt, jina_model, jina_tokenizer, device)
-    elif selected_model == "cross-encoder/stsb-roberta-base":
-        predicted_label, probabilities = cross_encoder_predict_relevance(system_prompt, user_prompt, stsb_model, stsb_tokenizer, device)
-    elif selected_model == "cross-encoder/ms-marco-MiniLM-L-6-v2":
-        predicted_label, probabilities = cross_encoder_predict_relevance(system_prompt, user_prompt, ms_model, ms_tokenizer, device)
-
-    probability_off_topic = probabilities[0][1] * 100
-    label = "Off-topic" if predicted_label==1 else "On-topic"
+def predict(system_prompt, user_prompt):
+    predicted_label_jina, probabilities_jina = embeddings_predict_relevance(system_prompt, user_prompt, jina_model, jina_tokenizer, device)
+    predicted_label_stsb, probabilities_stsb = cross_encoder_predict_relevance(system_prompt, user_prompt, stsb_model, stsb_tokenizer, device)
+    predicted_label_ms, probabilities_ms = cross_encoder_predict_relevance(system_prompt, user_prompt, ms_model, ms_tokenizer, device)
+
     result = f"""
-    **Prediction Summary**:
+    **Prediction Summary**
+
+    **1. Model: jinaai/jina-embeddings-v2-small-en**
+    - **Prediction**: {"🟥 Off-topic" if predicted_label_jina==1 else "🟩 On-topic"}
+    - **Probability of being off-topic**: {probabilities_jina[0][1]:.2%}
+
+    **2. Model: cross-encoder/stsb-roberta-base**
+    - **Prediction**: {"🟥 Off-topic" if predicted_label_stsb==1 else "🟩 On-topic"}
+    - **Probability of being off-topic**: {probabilities_stsb[0][1]:.2%}
 
-    - **Predicted Label**: {label}
-    - **Probability of Off-topic**: {probability_off_topic:.3f}%
+    **3. Model: cross-encoder/ms-marco-MiniLM-L-6-v2**
+    - **Prediction**: {"🟥 Off-topic" if predicted_label_ms==1 else "🟩 On-topic"}
+    - **Probability of being off-topic**: {probabilities_ms[0][1]:.2%}
     """
 
     return result
@@ -36,15 +40,8 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as app:
     gr.Markdown("# Off-Topic Classification using Fine-tuned Embeddings and Cross-Encoder Models")
 
     with gr.Row():
-        system_prompt = gr.Textbox(label="System Prompt")
-        user_prompt = gr.Textbox(label="User Prompt")
-
-    with gr.Row():
-        selected_model = gr.Dropdown(
-            ["jinaai/jina-embeddings-v2-small-en",
-             "cross-encoder/stsb-roberta-base",
-             "cross-encoder/ms-marco-MiniLM-L-6-v2"],
-            label="Select a model")
+        system_prompt = gr.TextArea(label="System Prompt", lines=5)
+        user_prompt = gr.TextArea(label="User Prompt", lines=5)
 
     # Button to run the prediction
     get_classfication = gr.Button("Check Content")
@@ -53,7 +50,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as app:
 
     get_classfication.click(
         fn=predict,
-        inputs=[system_prompt, user_prompt, selected_model],
+        inputs=[system_prompt, user_prompt],
         outputs=output_result
     )