File size: 3,064 Bytes
cdf4f6a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import gradio as gr
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("nickprock/AIRC_FAQ")
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() |