Update app.py

app.py

@@ -8,7 +8,7 @@ from sentence_transformers import SentenceTransformer
 def load_model(model_name):
     if model_name == "GLuCoSE-base-ja-v2":
         return SentenceTransformer("pkshatech/GLuCoSE-base-ja-v2")
-    elif model_name == "RoSEtta-base-ja":
+    elif model_name == "RoSEtta-base":
         return SentenceTransformer("pkshatech/RoSEtta-base", trust_remote_code=True)
     elif model_name == "ruri-large":
         return SentenceTransformer("cl-nagoya/ruri-large")
@@ -17,20 +17,21 @@ def load_model(model_name):
 def get_similarities(model_name, sentences):
     model = load_model(model_name)
 
-    if model_name == "ruri-large":
+    if model_name in ["GLuCoSE-base-ja-v2", "RoSEtta-base"]:
+        sentences = [
+            "query: " + s if i % 2 == 0 else "passage: " + s
+            for i, s in enumerate(sentences)
+        ]
+    elif model_name == "ruri-large":
         sentences = [
             "クエリ: " + s if i % 2 == 0 else "文章: " + s
             for i, s in enumerate(sentences)
         ]
 
     embeddings = model.encode(sentences, convert_to_tensor=True)
-
-    if model_name in ["GLuCoSE-base-ja-v2", "RoSEtta-base-ja"]:
-        similarities = model.similarity(embeddings, embeddings)
-    else:  # ruri-large
-        similarities = F.cosine_similarity(
-            embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2
-        )
+    similarities = F.cosine_similarity(
+        embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2
+    )
 
     return similarities.cpu().numpy()
 
@@ -45,7 +46,7 @@ def process_input(model_name, input_text):
     return format_similarities(similarities)
 
 
-models = ["GLuCoSE-base-ja-v2", "RoSEtta-base-ja", "ruri-large"]
+models = ["GLuCoSE-base-ja-v2", "RoSEtta-base", "ruri-large"]
 
 with gr.Blocks() as demo:
     gr.Markdown("# Sentence Similarity Demo")
@@ -55,7 +56,14 @@ with gr.Blocks() as demo:
             model_dropdown = gr.Dropdown(
                 choices=models, label="Select Model", value=models[0]
             )
-            input_text = gr.Textbox(lines=5, label="Input Sentences (one per line)")
+            input_text = gr.Textbox(
+                lines=5,
+                label="Input Sentences (one per line)",
+                placeholder="Enter query and passage pairs, alternating lines.",
+            )
+            gr.Markdown("""
+            **Note:** Prefixes ('query:' / 'passage:' or 'クエリ:' / '文章:') are added automatically. Just input your sentences.
+            """)
             submit_btn = gr.Button(value="Calculate Similarities")
 
         with gr.Column():
@@ -69,11 +77,11 @@ with gr.Blocks() as demo:
         examples=[
             [
                 "GLuCoSE-base-ja-v2",
-                "The weather is lovely today.\nIt's so sunny outside!\nHe drove to the stadium.",
+                "PKSHAはどんな会社ですか？\n研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。",
             ],
             [
-                "RoSEtta-base-ja",
-                "The weather is lovely today.\nIt's so sunny outside!\nHe drove to the stadium.",
+                "RoSEtta-base",
+                "PKSHAはどんな会社ですか？\n研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。",
             ],
             [
                 "ruri-large",