Update app.py

app.py

@@ -3,76 +3,107 @@ from sentence_transformers import SentenceTransformer
 import pandas as pd
 from datasets import load_dataset
 from annoy import AnnoyIndex
-import numpy as np
-
-# Load the dataset
-dataset = load_dataset("DeepMount00/CulturaViva-ITA")
-df = pd.DataFrame(dataset["train"])
-
-# Extract questions and answers
-questions = df["question"].tolist()
-answers = df["answer"].tolist()
-
-# Sentence-transformers models to test
-model_names = [
-    "nickprock/multi-sentence-BERTino",
-    "nickprock/sentence-bert-base-italian-uncased",
-    "nickprock/sentence-bert-base-italian-xxl-uncased",
-    "nickprock/mmarco-bert-base-italian-uncased",
-]
-
-models = {name: SentenceTransformer(name) for name in model_names}
-annoy_indexes = {}  # Store Annoy indexes for each model
-
-def build_annoy_index(model_name):
-    """Builds an Annoy index for a given model."""
-    model = models[model_name]
-    embeddings = model.encode(answers)
-    embedding_dim = embeddings.shape[1]
-    annoy_index = AnnoyIndex(embedding_dim, "angular")  # Use angular distance for cosine similarity
-    for i, embedding in enumerate(embeddings):
-        annoy_index.add_item(i, embedding)
-    annoy_index.build(10)  # Build with 10 trees
-    return annoy_index
-
-# Build Annoy indexes for each model
-for model_name in model_names:
-    annoy_indexes[model_name] = build_annoy_index(model_name)
-
-def find_similar_answer_annoy(question, model_name):
-    """Finds the most similar answer using Annoy."""
-    model = models[model_name]
-    annoy_index = annoy_indexes[model_name]
-    question_embedding = model.encode(question)
-    nearest_neighbors = annoy_index.get_nns_by_vector(question_embedding, 1)  # Get the nearest neighbor
-    best_answer_index = nearest_neighbors[0]
-    return answers[best_answer_index]
-
-def compare_models_annoy(question, model1_name, model2_name, model3_name, model4_name):
-    """Compares the results of different models using Annoy."""
-    answer1 = find_similar_answer_annoy(question, model1_name)
-    answer2 = find_similar_answer_annoy(question, model2_name)
-    answer3 = find_similar_answer_annoy(question, model3_name)
-    answer4 = find_similar_answer_annoy(question, model4_name)
-    return answer1, answer2, answer3, answer4
-
-iface = gr.Interface(
-    fn=compare_models_annoy,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Enter your question here..."),
-        gr.Dropdown(model_names, value=model_names[0], label="Model 1"),
-        gr.Dropdown(model_names, value=model_names[1], label="Model 2"),
-        gr.Dropdown(model_names, value=model_names[2], label="Model 3"),
-        gr.Dropdown(model_names, value=model_names[3], label="Model 4"),
-    ],
-    outputs=[
-        gr.Textbox(label=model_names[0]),
-        gr.Textbox(label=model_names[1]),
-        gr.Textbox(label=model_names[2]),
-        gr.Textbox(label=model_names[3]),
-    ],
-    title="Sentence Transformer Model Comparison (Annoy)",
-    description="Enter a question and compare the answers generated by different sentence-transformer models (using Annoy for faster search).",
-)
-
-iface.launch()
+import os
+
+try:
+    # Load the dataset (Italian subset, test split)
+    dataset = load_dataset("PhilipMay/stsb_multi_mt", name="it", split="test")
+    df = pd.DataFrame(dataset)
+
+    # Extract sentences (sentence1 and sentence2)
+    sentences1 = df["sentence1"].tolist()
+    sentences2 = df["sentence2"].tolist()
+
+    # Sentence-transformers models to test
+    model_names = [
+        "nickprock/multi-sentence-BERTino",
+        "nickprock/sentence-bert-base-italian-uncased",
+        "nickprock/sentence-bert-base-italian-xxl-uncased",
+        "nickprock/mmarco-bert-base-italian-uncased",
+    ]
+
+    models = {name: SentenceTransformer(name) for name in model_names}
+    annoy_indexes1 = {}  # Store Annoy indexes for sentence1
+    annoy_indexes2 = {}  # Store Annoy indexes for sentence2
+
+    def build_annoy_index(model_name, sentences):
+        """Builds an Annoy index for a given model and sentences."""
+        model = models[model_name]
+        embeddings = model.encode(sentences)
+        embedding_dim = embeddings.shape[1]
+        annoy_index = AnnoyIndex(embedding_dim, "angular")  # Use angular distance for cosine similarity
+        for i, embedding in enumerate(embeddings):
+            annoy_index.add_item(i, embedding)
+        annoy_index.build(10)  # Build with 10 trees
+        return annoy_index
+
+    # Build Annoy indexes for each model
+    for model_name in model_names:
+        annoy_indexes1[model_name] = build_annoy_index(model_name, sentences1)
+        annoy_indexes2[model_name] = build_annoy_index(model_name, sentences2)
+
+    def find_similar_sentence_annoy(sentence, model_name, sentence_list, annoy_index):
+        """Finds the most similar sentence using Annoy."""
+        model = models[model_name]
+        sentence_embedding = model.encode(sentence)
+        nearest_neighbors = annoy_index[model_name].get_nns_by_vector(sentence_embedding, 1)
+        best_sentence_index = nearest_neighbors[0]
+        return sentence_list[best_sentence_index]
+
+    def compare_models_annoy(sentence, model1_name, model2_name, model3_name, model4_name):
+        """Compares the results of different models using Annoy."""
+        sentence1_results = {}
+        sentence2_results = {}
+
+        sentence1_results[model1_name] = find_similar_sentence_annoy(sentence, model1_name, sentences1, annoy_indexes1)
+        sentence1_results[model2_name] = find_similar_sentence_annoy(sentence, model2_name, sentences1, annoy_indexes1)
+        sentence1_results[model3_name] = find_similar_sentence_annoy(sentence, model3_name, sentences1, annoy_indexes1)
+        sentence1_results[model4_name] = find_similar_sentence_annoy(sentence, model4_name, sentences1, annoy_indexes1)
+
+        sentence2_results[model1_name] = find_similar_sentence_annoy(sentence, model1_name, sentences2, annoy_indexes2)
+        sentence2_results[model2_name] = find_similar_sentence_annoy(sentence, model2_name, sentences2, annoy_indexes2)
+        sentence2_results[model3_name] = find_similar_sentence_annoy(sentence, model3_name, sentences2, annoy_indexes2)
+        sentence2_results[model4_name] = find_similar_sentence_annoy(sentence, model4_name, sentences2, annoy_indexes2)
+
+        return sentence1_results, sentence2_results
+
+    def format_results(sentence1_results, sentence2_results):
+        """Formats the results for display in Gradio."""
+        output_text = ""
+        for model_name in model_names:
+            output_text += f"**{model_name}**\n"
+            output_text += f"Most Similar Sentence from sentence1: {sentence1_results[model_name]}\n"
+            output_text += f"Most Similar Sentence from sentence2: {sentence2_results[model_name]}\n\n"
+        return output_text
+
+    def gradio_interface(sentence, model1_name, model2_name, model3_name, model4_name):
+        """Gradio interface function."""
+        sentence1_results, sentence2_results = compare_models_annoy(sentence, model1_name, model2_name, model3_name, model4_name)
+        return format_results(sentence1_results, sentence2_results)
+
+    iface = gr.Interface(
+        fn=gradio_interface,
+        inputs=[
+            gr.Textbox(lines=2, placeholder="Enter your sentence here..."),
+            gr.Dropdown(model_names, value=model_names[0], label="Model 1"),
+            gr.Dropdown(model_names, value=model_names[1], label="Model 2"),
+            gr.Dropdown(model_names, value=model_names[2], label="Model 3"),
+            gr.Dropdown(model_names, value=model_names[3], label="Model 4"),
+        ],
+        outputs=gr.Markdown(),
+        title="Sentence Transformer Model Comparison (Annoy)",
+        description="Enter a sentence and compare the most similar sentences generated by different sentence-transformer models (using Annoy for faster search) from both sentence1 and sentence2.",
+    )
+
+    iface.launch()
+
+except Exception as e:
+    print(f"Error loading dataset: {e}")
+    iface = gr.Interface(
+        fn=lambda: "Dataset loading failed. Check console for details.",
+        inputs=[],
+        outputs=gr.Textbox(),
+        title="Dataset Loading Error",
+        description="There was an error loading the dataset.",
+    )
+    iface.launch()