Update app.py

app.py

@@ -1,93 +1,50 @@
-import gradio as gr
 from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
-import multiprocessing
-import chromadb
-import hashlib
-
-# Carga el modelo
-model = SentenceTransformer('Maite89/Roberta_finetuning_semantic_similarity_stsb_multi_mt')
-
-# Crea el cliente ChromaDB
-chroma_client = chromadb.Client()
-collection = chroma_client.create_collection(name="my_collection")
-
-def generate_hash(text):
-    return hashlib.md5(text.encode('utf-8')).hexdigest()
-
-# Función para obtener embeddings del modelo
-import sqlite3
+from accelerate import Accelerator
 import gradio as gr
-from sentence_transformers import SentenceTransformer
-from sklearn.metrics.pairwise import cosine_similarity
-import numpy as np
-import multiprocessing
 
-# Inicializa la base de datos y crea la tabla si no existe
-conn = sqlite3.connect('embeddings.db')
-c = conn.cursor()
-c.execute('''CREATE TABLE IF NOT EXISTS embeddings
-             (sentence TEXT PRIMARY KEY, embedding BLOB)''')
-conn.commit()
+# Inicializa el Accelerator
+accelerator = Accelerator()
 
-# Carga el modelo
+# 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):
-    # Intenta recuperar los embeddings de la base de datos
-    embeddings = []
-    new_sentences = []
-    for sentence in sentences:
-        c.execute('SELECT embedding FROM embeddings WHERE sentence=?', (sentence,))
-        result = c.fetchone()
-        if result:
-            embeddings.append(np.frombuffer(result[0], dtype=np.float32))
-        else:
-            new_sentences.append(sentence)
-    
-    # Si hay nuevas sentencias, obtén los embeddings y almacénalos en la base de datos
-    if new_sentences:
-        new_embeddings = model.encode(new_sentences, show_progress_bar=False)
-        embeddings.extend(new_embeddings)
-        c.executemany('INSERT INTO embeddings VALUES (?,?)',
-                      [(sent, emb.tobytes()) for sent, emb in zip(new_sentences, new_embeddings)])
-        conn.commit()
-    
-    return embeddings
+    # 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 comparar las sentencias
-def calculate_similarity(args):
-    source_embedding, compare_embedding = args
-    return cosine_similarity(source_embedding.reshape(1, -1), compare_embedding.reshape(1, -1))[0][0]
+# 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("--")
     
-    # Obtiene todos los embeddings a la vez para acelerar el proceso
+    # Obtener todos los embeddings de una vez para acelerar el proceso
     all_sentences = [source_sentence] + compare_list
     all_embeddings = get_embeddings(all_sentences)
     
-    # Prepara los datos para el multiprocesamiento
+    # No se necesita multiprocesamiento si usamos Accelerate ya que esto se maneja internamente
     source_embedding = all_embeddings[0]
-    data_for_multiprocessing = [(source_embedding, emb) for emb in all_embeddings[1:]]
-    
-    # Utiliza un pool de procesos para calcular las similitudes en paralelo
-    with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
-        similarities = pool.map(calculate_similarity, data_for_multiprocessing)
+    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
+# Definir las interfaces de entrada y salida para Gradio
 iface = gr.Interface(
-    fn=compare, 
-    inputs=["text", "text"], 
-    outputs="text",
-    live=False  
+    fn=compare,
+    inputs=[gr.inputs.Textbox(lines=2, placeholder="Enter source sentence here..."),
+            gr.inputs.Textbox(lines=10, placeholder="Enter sentences to compare, separated by '--'...")],
+    outputs=gr.outputs.Textbox(),
+    live=False
 )
 
-# Inicia la interfaz de Gradio
+# Iniciar la interfaz de Gradio
 iface.launch()
-conn.close()