update sts

README.md

@@ -4,7 +4,7 @@ emoji: 🏢
 colorFrom: green
 colorTo: gray
 sdk: gradio
-sdk_version: 5.11.0
+sdk_version: 5.23.1
 app_file: app.py
 pinned: false
 ---

app.py

@@ -73,6 +73,9 @@ with gr.Blocks() as demo:
 
         with gr.Row():
             with gr.Column():
+                model = gr.Dropdown(
+                    ["Lajavaness/bilingual-embedding-large", "sentence-transformers/all-mpnet-base-v2",
+                     "intfloat/multilingual-e5-large-instruct"], label="model", interactive=True)
                 score_sts = gr.Number(label="score", value=0.96, interactive=True)
                 submit_button_sts = gr.Button("Submit", variant="primary")
 
@@ -85,7 +88,7 @@ with gr.Blocks() as demo:
 
                 submit_button_sts.click(
                     fn=sts,
-                    inputs=[upload_button_sentences1, upload_button_sentences2, score_sts],
+                    inputs=[model, upload_button_sentences1, upload_button_sentences2, score_sts],
                     outputs=output_sts
                 )
 

sts.py

@@ -5,20 +5,19 @@ import torch
 from datasets import Dataset
 from sentence_transformers import SentenceTransformer
 
-def sts(data1, data2, score):
+
+def sts(modelname, data1, data2, score):
     st = time.time()
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = SentenceTransformer(
-        "sentence-transformers/all-MiniLM-L6-v2",
-        backend="openvino",
-        model_kwargs={"file_name": "openvino/openvino_model.xml"},
+        modelname,
         device=device,
         trust_remote_code=True,
     )
 
-    sentences1 = Dataset.from_pandas(pd.read_csv(data1, on_bad_lines='skip', header=0, names=["text"], sep="\t"))
-    sentences2 = Dataset.from_pandas(pd.read_csv(data2, on_bad_lines='skip', header=0, names=["text"], sep="\t"))
+    sentences1 = Dataset.from_pandas(pd.read_csv(data1, on_bad_lines='skip', header=0, sep="\t"))
+    sentences2 = Dataset.from_pandas(pd.read_csv(data2, on_bad_lines='skip', header=0, sep="\t"))
 
     embeddings1 = model.encode(sentences1["text"], normalize_embeddings=True, batch_size=1024,
                                show_progress_bar=True)
@@ -45,12 +44,12 @@ def sts(data1, data2, score):
                .join(df_sentences2, on="column_index"))
 
     df_long = df_long.rename({
-        "text": "setences1",
+        "text": "sentences1",
         "text_right": "sentences2",
-    })
+    }).drop(["row_index", "column_index"])
 
     elapsed_time = time.time() - st
     print('Execution time:', time.strftime("%H:%M:%S", time.gmtime(elapsed_time)))
 
-    return df_long.select(["score", "setences1", "sentences2"]).filter(pl.col("score") > score).sort(["score"],
-                                                                                                    descending=True)
+    return df_long.filter(pl.col("score") > score).sort(["score"],
+                                                        descending=True)

utils.py

@@ -6,7 +6,7 @@ import polars as pl
 import time
 
 def getDataFrame(path):
-    data = pd.read_csv(path, on_bad_lines='skip', header=0, names=["text"], sep="\t")
+    data = pd.read_csv(path, on_bad_lines='skip', header=0, sep="\t")
     return pl.from_pandas(data)
 
 def save_to_csv(dataframe):