Similacion / app.py
antagonico's picture
Update app.py
302c737
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
from accelerate import Accelerator
import gradio as gr
# Inicializa el Accelerator
accelerator = Accelerator()
# Cargar el modelo y colocarlo en el dispositivo adecuado
model = SentenceTransformer('Maite89/Roberta_finetuning_semantic_similarity_stsb_multi_mt')
model, _ = accelerator.prepare(model, model)
# Funci贸n para obtener embeddings del modelo
def get_embeddings(sentences):
# Preparar los datos para ejecuci贸n acelerada
sentences = accelerator.prepare(sentences)
return model.encode(sentences, show_progress_bar=False, convert_to_tensor=True)
# Funci贸n para calcular la similitud
def calculate_similarity(arguments):
source_embedding, compare_embedding = arguments
return cosine_similarity([source_embedding], [compare_embedding])[0][0]
# Funci贸n para comparar oraciones
def compare(source_sentence, compare_sentences):
compare_list = compare_sentences.split("--")
# Obtener todos los embeddings de una vez para acelerar el proceso
all_sentences = [source_sentence] + compare_list
all_embeddings = get_embeddings(all_sentences)
# No se necesita multiprocesamiento si usamos Accelerate ya que esto se maneja internamente
source_embedding = all_embeddings[0]
similarities = [calculate_similarity((source_embedding, emb)) for emb in all_embeddings[1:]]
return ', '.join([str(sim) for sim in similarities])
# Define las interfaces de entrada y salida de Gradio
iface = gr.Interface(
fn=compare,
inputs=["text", "text"],
outputs="text",
live=False
)
# Iniciar la interfaz de Gradio
iface.launch()