import pandas as pd
from ast import literal_eval
from tqdm import tqdm
tqdm.pandas()

df = pd.read_csv('/mnt/home/dongkeun/L2U/outputs/doc_id_reset_count_generation.csv')
df['prefix'] = df['prefix'].apply(literal_eval)
df['suffix'] = df['suffix'].apply(literal_eval)
df['wrong_generation'] = df['wrong_generation'].apply(literal_eval)

df['right_generation'] = df.apply(lambda x: x['prefix'] + x['suffix'], axis=1)

df['same'] = df.apply(lambda x: x['right_generation'] == x['wrong_generation'], axis=1)
df = df[df['same']]
df = df[~df['eoe_init'].isin([0.0, 127.0])]
print(df.shape)
print(df['text'])
df.to_csv('outputs/tmp.csv', index=False)
