feat: 添加不同query对比

app.py

@@ -13,65 +13,35 @@ model4 = SentenceTransformer("aspire/acge_text_embedding")
 model5 = SentenceTransformer("intfloat/multilingual-e5-large")
 
 @spaces.GPU
-def generate(input1, input2):
-    if len(input1) < 1:
-        input1 = "How do I access the index while iterating over a sequence with a for loop?"
-    if len(input2) < 1:
-        input2 = "# Use the built-in enumerator\nfor idx, x in enumerate(xs):\n    print(idx, x)"
-
-    embeddings1 = model1.encode(
-        [
-            input1,
-            input2,
-        ]
-    )
-    score1 = cos_sim(embeddings1[0], embeddings1[1])
-
-    embeddings2 = model2.encode(
-        [
-            input1,
-            input2,
-        ]
-    )
-    score2 = cos_sim(embeddings2[0], embeddings2[1])
-
-    embeddings3 = model3.encode(
-        [
-            input1,
-            input2,
-        ]
-    )
-    score3 = cos_sim(embeddings3[0], embeddings3[1])
-
-    embeddings4 = model4.encode(
-        [
-            input1,
-            input2,
-        ]
-    )
-    score4 = cos_sim(embeddings4[0], embeddings4[1])
-
-    embeddings5 = model5.encode(
-        [
-            input1,
-            input2,
-        ]
-    )
-    score5 = cos_sim(embeddings5[0], embeddings5[1])
-
-    return score1, score2, score3, score4, score5
+def generate(query1, query2, source_code):
+    if len(query1) < 1:
+        query1 = "How do I access the index while iterating over a sequence with a for loop?"
+    if len(query2) < 1:
+        query2 = "get a list of all the keys in a dictionary"
+    if len(source_code) < 1:
+        source_code = "# Use the built-in enumerator\nfor idx, x in enumerate(xs):\n    print(idx, x)"
+
+    results = []
+    for model in [model1, model2, model3, model4, model5]:
+        embeddings = model.encode([query1, query2, source_code])
+        score1 = cos_sim(embeddings[0], embeddings[2])
+        score2 = cos_sim(embeddings[1], embeddings[2])
+        results.append((float(score1), float(score2)))
+
+    return results
 
 gr.Interface(
     fn=generate,
     inputs=[
-        gr.Text(label="input1", placeholder="How do I access the index while iterating over a sequence with a for loop?"),
-        gr.Text(label="input2", placeholder="# Use the built-in enumerator\nfor idx, x in enumerate(xs):\n    print(idx, x)"),
+        gr.Text(label="query1", placeholder="How do I access the index while iterating over a sequence with a for loop?"),
+        gr.Text(label="query2", placeholder="get a list of all the keys in a dictionary"),
+        gr.Text(label="code", placeholder="# Use the built-in enumerator\nfor idx, x in enumerate(xs):\n    print(idx, x)"),
     ],
     outputs=[
-        gr.Text(label="jinaai/jina-embeddings-v2-base-code"),
-        gr.Text(label="jinaai/jina-embeddings-v2-base-en"),
-        gr.Text(label="jinaai/jina-embeddings-v2-base-zh"),
-        gr.Text(label="aspire/acge_text_embedding"),
-        gr.Text(label="intfloat/multilingual-e5-large"),
+        gr.Dataframe(
+            headers=["Query1 Score", "Query2 Score"],
+            label="Similarity Scores",
+            row_labels=["jinaai/jina-embeddings-v2-base-code", "jinaai/jina-embeddings-v2-base-en", "jinaai/jina-embeddings-v2-base-zh", "aspire/acge_text_embedding", "intfloat/multilingual-e5-large"]
+        )
     ],
 ).launch()