update

.gitignore

@@ -1,5 +1,6 @@
 .idea
+flagged
 data/example
 data/tmp
 
-output.csv
+output.csv

utils/pipeline.py

@@ -7,6 +7,12 @@ from utils.models import SBert
 
 
 def p0_originality(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
+    """
+    row-wise
+    :param df:
+    :param model_name:
+    :return:
+    """
     assert 'prompt' in df.columns
     assert 'response' in df.columns
     model = SBert(model_name)
@@ -22,12 +28,18 @@ def p0_originality(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
 
 
 def p1_flexibility(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
+    """
+    group-wise
+    :param df:
+    :param model_name:
+    :return:
+    """
     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:
+    def get_flexibility(responses: List[str]) -> float:
         responses_vec = [model(_) for _ in responses]
         count = 0
         score = 0
@@ -40,7 +52,7 @@ def p1_flexibility(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
         return score / count
 
     df_out = df.groupby(by=['id', 'prompt']) \
-        .agg({'id': 'first', 'prompt': 'first', 'response': get_cos_sim}) \
+        .agg({'id': 'first', 'prompt': 'first', 'response': get_flexibility}) \
         .rename(columns={'response': 'flexibility'}) \
         .reset_index(drop=True)
     return df_out