|
import pandas as pd |
|
from sentence_transformers.util import cos_sim |
|
|
|
from utils.models import SBert |
|
|
|
|
|
def get_cos_sim(model, prompt: str, response: str) -> float: |
|
prompt_vec = model(prompt) |
|
response_vec = model(response) |
|
score = cos_sim(prompt_vec, response_vec).item() |
|
return score |
|
|
|
|
|
def batch_cos_sim(df: pd.DataFrame, model_name) -> pd.DataFrame: |
|
|
|
assert 'prompt' in df.columns |
|
assert 'response' in df.columns |
|
model = SBert(model_name) |
|
df['originality'] = df.apply(lambda x: 1 - get_cos_sim(model, x['prompt'], x['response']), axis=1) |
|
return df |
|
|
|
|
|
if __name__ == '__main__': |
|
_df = pd.read_csv('data/example_1.csv') |
|
_df_o = batch_cos_sim(_df, 'paraphrase-multilingual-MiniLM-L12-v2') |
|
|
