update app.py

app.py

@@ -8,50 +8,51 @@ model = SentenceTransformer("Omartificial-Intelligence-Space/Arabic-Nli-Matryosh
 # Define the labels for NLI
 labels = ["contradiction", "entailment", "neutral"]
 
-# Define the Matryoshka dimension
-matryoshka_dim = 64
-
 # Function to compute similarity and classify relationship
-def predict(sentence1, sentence2):
-    sentences = [sentence1, sentence2]
-    embeddings = model.encode(sentences)
-    
-    # Shrink the embedding dimensions
-    embeddings = embeddings[..., :matryoshka_dim]
-    
-    # Compute cosine similarity between the two sentences
-    similarity_score = util.cos_sim(embeddings[0], embeddings[1])
+def predict(mode, sentence1, sentence2=None, sentence3=None, sentence4=None, dimension="64"):
+    dimension = int(dimension)
+    if mode == "Compare one to three":
+        if sentence2 is None or sentence3 is None or sentence4 is None:
+            return "Please provide three sentences for comparison.", {}
+        sentences = [sentence1, sentence2, sentence3, sentence4]
+    else:
+        if sentence2 is None:
+            return "Please provide the second sentence for comparison.", {}
+        sentences = [sentence1, sentence2]
     
-    # Placeholder logic for NLI (needs to be replaced with actual model inference)
-    # This is just an example; in reality, you need a classifier trained for NLI
-    scores = np.random.rand(3)  # Replace this with actual model prediction logic
-    scores = scores / scores.sum()  # Normalize to sum to 1
+    embeddings = model.encode(sentences)
+    embeddings = embeddings[..., :dimension]
     
-    label_probs = {labels[i]: float(scores[i]) for i in range(len(labels))}
+    if mode == "Compare one to three":
+        similarities = util.cos_sim(embeddings[0], embeddings[1:])
+        similarity_scores = {f"Sentence {i+2}": float(similarities[0, i]) for i in range(3)}
+    else:
+        similarity_score = util.cos_sim(embeddings[0], embeddings[1])
+        similarity_scores = {"Similarity Score": float(similarity_score)}
     
-    return similarity_score.item(), label_probs
+    return similarity_scores
 
 # Define inputs and outputs for Gradio interface
-inputs = [
-    gr.Textbox(lines=2, placeholder="Enter the first sentence here...", label="Sentence 1"),
-    gr.Textbox(lines=2, placeholder="Enter the second sentence here...", label="Sentence 2")
-]
+mode_dropdown = gr.Dropdown(choices=["Compare two sentences", "Compare one to three"], label="Mode", default="Compare two sentences")
+dimension_dropdown = gr.inputs.Dropdown(choices=["768", "512", "256", "128", "64"], label="Embedding Dimension", default="64")
+sentence1_input = gr.Textbox(lines=2, placeholder="Enter the first sentence here...", label="Sentence 1")
+sentence2_input = gr.Textbox(lines=2, placeholder="Enter the second sentence here...", label="Sentence 2 (or first of three for mode)", optional=True)
+sentence3_input = gr.Textbox(lines=2, placeholder="Enter the third sentence here...", label="Sentence 3", optional=True)
+sentence4_input = gr.Textbox(lines=2, placeholder="Enter the fourth sentence here...", label="Sentence 4", optional=True)
 
-outputs = [
-    gr.Textbox(label="Similarity Score"),
-    gr.Label(num_top_classes=3, label="Label Probabilities")
-]
+inputs = [mode_dropdown, sentence1_input, sentence2_input, sentence3_input, sentence4_input, dimension_dropdown]
+outputs = gr.JSON(label="Similarity Scores")
 
 examples = [
-    ["يجلس شاب ذو شعر أشقر على الحائط يقرأ جريدة بينما تمر امرأة وفتاة شابة.", "ذكر شاب ينظر إلى جريدة بينما تمر إمرأتان بجانبه"],
-    ["الشاب نائم بينما الأم تقود ابنتها إلى الحديقة", "ذكر شاب ينظر إلى جريدة بينما تمر إمرأتان بجانبه"]
+    ["Compare one to three", "يجلس شاب ذو شعر أشقر على الحائط يقرأ جريدة بينما تمر امرأة وفتاة شابة.", "ذكر شاب ينظر إلى جريدة بينما تمر إمرأتان بجانبه", "الشاب نائم بينما الأم تقود ابنتها إلى الحديقة", "رجل يقرأ الجريدة في الحديقة", "64"],
+    ["Compare two sentences", "يجلس شاب ذو شعر أشقر على الحائط يقرأ جريدة بينما تمر امرأة وفتاة شابة.", "ذكر شاب ينظر إلى جريدة بينما تمر إمرأتان بجانبه", None, None, "64"]
 ]
 
 # Create Gradio interface
 gr.Interface(
     fn=predict,
-    title="Arabic Sentence Similarity and NLI Classification",
-    description="Compute the semantic similarity and classify the relationship between two Arabic sentences using a SentenceTransformer model.",
+    title="Arabic Sentence Similarity with Matryoshka Model",
+    description="Compute the semantic similarity between Arabic sentences using the Matryoshka SentenceTransformer model.",
     inputs=inputs,
     examples=examples,
     outputs=outputs,