update mining

app.py

@@ -25,6 +25,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_mining = gr.Number(label="score", value=0.96, interactive=True)
                 submit_button_mining = gr.Button("Submit", variant="primary")
 
@@ -35,7 +38,7 @@ with gr.Blocks() as demo:
 
                 submit_button_mining.click(
                     fn=mining,
-                    inputs=[upload_button_sentences, score_mining],
+                    inputs=[model, upload_button_sentences, score_mining],
                     outputs=output_mining
                 )
 

mining.py

@@ -6,15 +6,15 @@ from datasets import Dataset
 from sentence_transformers import SentenceTransformer
 from sentence_transformers.util import paraphrase_mining
 
-def mining(path, score):
+
+def mining(modelname, path, score):
     st = time.time()
-    data = Dataset.from_pandas(pd.read_csv(path, on_bad_lines='skip', header=0, names=["text"], sep="\t"))
+    data = Dataset.from_pandas(pd.read_csv(path, on_bad_lines='skip', header=0, sep="\t"))
+    original_df = pd.read_csv(path, on_bad_lines='skip', header=0, sep="\t")
 
     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,
     )
@@ -34,13 +34,24 @@ def mining(path, score):
 
     union_df = pl.DataFrame(data.to_pandas())
 
+    original_columns = original_df.columns.tolist()
+
+    additional_cols = []
+    for col in original_columns:
+        if col != "text":
+            additional_cols.extend([
+                union_df.select(pl.col(col)).to_series()[df["sentence_1"].cast(pl.Int32)].alias(f"{col}_1"),
+                union_df.select(pl.col(col)).to_series()[df["sentence_2"].cast(pl.Int32)].alias(f"{col}_2")
+            ])
+
     df = df.with_columns([
         pl.col("score").round(3).cast(pl.Float32),
         union_df.select(pl.col("text")).to_series()[df["sentence_1"].cast(pl.Int32)].alias("sentence_1"),
         union_df.select(pl.col("text")).to_series()[df["sentence_2"].cast(pl.Int32)].alias("sentence_2"),
+        *additional_cols
     ]).filter(pl.col("score") > score).sort(["score"], descending=True)
 
     elapsed_time = time.time() - st
     print('Execution time:', time.strftime("%H:%M:%S", time.gmtime(elapsed_time)))
 
-    return df
+    return df