tcyang/TransDis-CreativityAutoAssessment-V2

1from typing import List2 3import pandas as pd4from sentence_transformers.util import cos_sim5 6from utils.models import ModelWithPooling7 8 9def p0_originality(df: pd.DataFrame, model_name: str, pooling: str) -> pd.DataFrame:10    """11    row-wise12    :param df:13    :param model_name:14    :return:15    """16    assert 'prompt' in df.columns17    assert 'response' in df.columns18    model = ModelWithPooling(model_name)19 20    def get_cos_sim(prompt: str, response: str) -> float:21        prompt_vec = model(text=prompt, pooling=pooling)22        response_vec = model(text=response, pooling=pooling)23        score = cos_sim(prompt_vec, response_vec).item()24        return score25 26    df['originality'] = df.apply(lambda x: 1 - get_cos_sim(x['prompt'], x['response']), axis=1)27    return df28 29 30def p1_flexibility(df: pd.DataFrame, model_name: str, pooling: str) -> pd.DataFrame:31    """32    group-wise33    :param df:34    :param model_name:35    :return:36    """37    assert 'prompt' in df.columns38    assert 'response' in df.columns39    assert 'id' in df.columns40    model = ModelWithPooling(model_name)41 42    def get_flexibility(responses: List[str]) -> float:43        responses_vec = [model(text=_, pooling=pooling) for _ in responses]44        score = 045        for i in range(len(responses_vec) - 1):46            score += 1 - cos_sim(responses_vec[i], responses_vec[i + 1]).item()47        return score48 49    df_out = df.groupby(by=['id', 'prompt']) \50        .agg({'id': 'first', 'prompt': 'first', 'response': get_flexibility}) \51        .rename(columns={'response': 'flexibility'}) \52        .reset_index(drop=True)53    return df_out54 55 56if __name__ == '__main__':57    _df_input = pd.read_csv('data/tmp/example_3.csv')58    _df_0 = p0_originality(_df_input, 'paraphrase-multilingual-MiniLM-L12-v2')59    _df_1 = p1_flexibility(_df_input, 'paraphrase-multilingual-MiniLM-L12-v2')60