mirror of
https://github.com/wassname/discovering_latent_knowledge.git
synced 2026-10-30 13:10:08 +08:00
43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
from datasets import load_dataset, load_from_disk
|
|
from src.helpers.typing import int16_to_float, float_to_int16
|
|
|
|
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 or (x.ndim==1 and len(x)==1)):
|
|
row[k]=x.item()
|
|
if isinstance(x, list) and len(x)==1:
|
|
row[k]=x[0]
|
|
return row
|
|
|
|
|
|
def ds2df(ds, cols=None):
|
|
"""one of our custom datasets into a dataframe
|
|
|
|
dropping the large arrays and lists"""
|
|
if cols is None:
|
|
r = ds[0]
|
|
# get all the columns that not large lists or arrays
|
|
cols = [k for k,v in r.items() if (isinstance(v, np.ndarray) and v.size<3) or not isinstance(v, (list, np.ndarray))]
|
|
|
|
df = ds.select_columns(cols)
|
|
df = pd.DataFrame([rows_item(r) for r in df])
|
|
|
|
# derived
|
|
df['dir_true'] = df['ans0']
|
|
# df['conf'] = (df['ans0']).abs9)
|
|
# df['llm_prob'] = (df['ans0']+df['ans1'])/2
|
|
df['llm_ans'] = df['ans0']>0.5
|
|
# df['desired_ans'] = df.label ^ df.lie
|
|
return df
|
|
|
|
def load_ds(f):
|
|
ds = load_from_disk(f)
|
|
ks = ds['large_arrays_keys'][0]
|
|
return ds.map(lambda x: {k: int16_to_float(torch.from_numpy(ds[k])) for k in ks})
|