Files
discovering_latent_knowledge/src/datasets/load.py
T
2023-07-28 18:29:48 +08:00

29 lines
795 B
Python

import numpy as np
import pandas as pd
def rows_item(row):
"""
transform a row by turning singe dim arrays into items
"""
for k,x in row.items():
if isinstance(x, np.ndarray) and x.ndim==0:
row[k]=x.item()
return row
def ds_info2df(ds):
info = list(ds['info'])
d = pd.DataFrame([rows_item(r) for r in info])
return d
def ds2df(ds):
df = ds_info2df(ds)
df_ans = ds.select_columns(['ans1', 'ans2', 'true', 'index', 'prob_y', 'prob_n', 'version']).with_format("numpy").to_pandas()
df = pd.concat([df, df_ans], axis=1)
# derived
df['dir_true'] = df['ans2'] - df['ans1']
df['conf'] = (df['ans1']-df['ans2']).abs()
df['llm_prob'] = (df['ans1']+df['ans2'])/2
df['llm_ans'] = df['llm_prob']>0.5
return df