update p2_flexibility, ui

.gitignore

@@ -1,3 +1,5 @@
-data
 .idea
-*.csv
+data/example
+data/tmp
+
+output.csv

app.py

@@ -50,18 +50,13 @@ def process(task_name: str,
         return {'Error': e}, None, None
 
 
+# input
 task_name_dropdown = gr.components.Dropdown(
     label='Task Name',
     value='Originality',
     choices=['Originality', 'Flexibility']
 )
 
-model_name_input = gr.components.Textbox(
-    value='paraphrase-multilingual-MiniLM-L12-v2',
-    lines=1,
-    type='text'
-)
-
 model_name_dropdown = gr.components.Dropdown(
     label='Model Name',
     value=list_models[0],
@@ -69,11 +64,16 @@ model_name_dropdown = gr.components.Dropdown(
 )
 
 text_input = gr.components.Textbox(
-    value='id,prompt,response\n',
+    value=open('data/example.csv', 'r').read(),
     lines=10,
     type='text'
 )
 
+# output
+file_input = gr.components.File(label='Input File',
+                                file_count='single',
+                                file_types=['', '.', '.csv', '.xls', '.xlsx'])
+
 text_output = gr.components.Textbox(
     label='Output',
     type='text'
@@ -83,16 +83,14 @@ dataframe_output = gr.components.Dataframe(
     label='DataFrame'
 )
 
-description = open('description.txt', 'r').read()
-
 file_output = gr.components.File(label='Output File',
                                  file_count='single',
                                  file_types=['', '.', '.csv', '.xls', '.xlsx'])
 
 app = gr.Interface(
     fn=process,
-    inputs=[task_name_dropdown, model_name_dropdown, text_input, 'file'],
+    inputs=[task_name_dropdown, model_name_dropdown, text_input, file_input],
     outputs=[text_output, dataframe_output, file_output],
-    description=description
+    description=open('data/description.txt', 'r').read()
 )
 app.launch()

data/example.csv

@@ -0,0 +1,10 @@
+id,prompt,response
+1,床单,过滤器
+1,床单,做成渔网捞鱼
+1,床单,做成枕头
+1,牙刷,捅人
+1,牙刷,用作鞋拔
+1,牙刷,当飞镖扔
+2,床单,做被子
+2,床单,保暖
+2,床单,绑在树上做成吊床

utils/models.py

@@ -1,6 +1,6 @@
-import numpy as np
-import torch
 from functools import lru_cache
+
+import torch
 from sentence_transformers import SentenceTransformer
 
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -20,6 +20,6 @@ class SBert:
         self.model = SentenceTransformer(path, device=DEVICE)
 
     @lru_cache(maxsize=10000)
-    def __call__(self, x) -> np.ndarray:
-        y = self.model.encode(x)
+    def __call__(self, x) -> torch.Tensor:
+        y = self.model.encode(x, convert_to_tensor=True)
         return y

utils/pipeline.py

@@ -9,22 +9,37 @@ def p0_originality(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
     assert 'response' in df.columns
     model = SBert(model_name)
 
-    def get_cos_sim(model, prompt: str, response: str) -> float:
+    def get_cos_sim(prompt: str, response: str) -> float:
         prompt_vec = model(prompt)
         response_vec = model(response)
         score = cos_sim(prompt_vec, response_vec).item()
         return score
 
-    df['originality'] = df.apply(lambda x: 1 - get_cos_sim(model, x['prompt'], x['response']), axis=1)
+    df['originality'] = df.apply(lambda x: 1 - get_cos_sim(x['prompt'], x['response']), axis=1)
     return df
 
 
 def p1_flexibility(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
-    df = p0_originality(df, model_name)
+    assert 'prompt' in df.columns
+    assert 'response' in df.columns
     assert 'id' in df.columns
+    model = SBert(model_name)
+
+    def get_cos_sim(responses: list[str]) -> float:
+        responses_vec = [model(_) for _ in responses]
+        count = 0
+        score = 0
+        for i in range(len(responses_vec)):
+            for j in range(1, len(responses_vec)):
+                if i == j:
+                    continue
+                score += cos_sim(responses_vec[i], responses_vec[j]).item()
+                count += 1
+        return score / count
+
     df_out = df.groupby(by=['id', 'prompt']) \
-        .agg({'id': 'first', 'prompt': 'first', 'originality': 'mean'}) \
-        .rename(columns={'originality': 'flexibility'}) \
+        .agg({'id': 'first', 'prompt': 'first', 'response': get_cos_sim}) \
+        .rename(columns={'response': 'flexibility'}) \
         .reset_index(drop=True)
     return df_out