data/helpsteer.py

import json,gzip
from pathlib import Path
import pandas as pd,numpy as np
from data.splits import prompt_id,response_id,assign_split,assert_disjoint

CRITERIA=['correctness','coherence','complexity','verbosity']
def main():
    info=json.loads(Path('outputs/downloads/data.json').read_text()); root=Path(info['path']); out=Path('data/processed'); out.mkdir(exist_ok=True)
    read=lambda name:[json.loads(x) for x in gzip.open(root/name,'rt')]
    rawtrain=read('train.jsonl.gz'); rawtest=read('validation.jsonl.gz'); prefs=read('preference/preference.jsonl.gz')
    testids={prompt_id(x['prompt']) for x in rawtest}; trainids={prompt_id(x['prompt']) for x in rawtrain}; overlap=testids&trainids
    examples={}; labels=[]
    def add(prompt,response):
        rid=response_id(prompt,response); pid=prompt_id(prompt)
        examples[rid]={'example_id':rid,'response_id':rid,'prompt_id':pid,'prompt':prompt,'response':response,'split':assign_split(pid,testids)}
        return rid
    for r in rawtrain+rawtest:
        rid=add(r['prompt'],r['response']); labels.append({'response_id':rid,**{k:float(r[k]) for k in CRITERIA+['helpfulness']}})
    pairs=[]
    for p in prefs:
        a=add(p['prompt'],p['response_1']); b=add(p['prompt'],p['response_2']); strength=float(p['preference_strength']); pid=prompt_id(p['prompt'])
        pairs.append({'pair_id':f'{a}_{b}','prompt_id':pid,'a':a,'b':b,'label':1. if strength<0 else 0. if strength>0 else .5,'strength':strength,'tie':strength==0,'official_split':p['split'],'split':assign_split(pid,testids)})
    ex=pd.DataFrame(examples.values()).sort_values('response_id').reset_index(drop=True); ex['feature_index']=np.arange(len(ex)); assert_disjoint(ex)
    lab=pd.DataFrame(labels).groupby('response_id',as_index=False).mean(); pair=pd.DataFrame(pairs).drop_duplicates('pair_id'); pair=pair[pair.a!=pair.b].copy()
    lookup=ex.set_index('response_id').feature_index; pair['a_index']=pair.a.map(lookup); pair['b_index']=pair.b.map(lookup)
    assert not pair[['a_index','b_index']].isna().any().any()
    assert all(pair['split'].to_numpy()==ex.set_index('response_id').loc[pair.a,'split'].to_numpy())
    assert all(pair['split'].to_numpy()==ex.set_index('response_id').loc[pair.b,'split'].to_numpy())
    ex.to_parquet(out/'responses.parquet',index=False); lab.to_parquet(out/'attributes.parquet',index=False); pair.to_parquet(out/'preferences.parquet',index=False)
    joined=ex.merge(lab,on='response_id',how='left'); rows=[]
    for s,g in ex.groupby('split'):
        p=pair[pair.split==s]; rows.append({'split':s,'prompts':g.prompt_id.nunique(),'responses':len(g),'preference_pairs':len(p),'non_tie_pairs':sum(~p.tie),'ties':sum(p.tie),'rated_responses':joined[joined.split==s].correctness.notna().sum()})
    stats=pd.DataFrame(rows); chars=ex.response.str.len().describe(percentiles=[.5,.75,.9,.95,.99]); attrs=joined[CRITERIA+['helpfulness']]
    report='# Dataset report\n\nSource: [NVIDIA HelpSteer2](https://huggingface.co/datasets/nvidia/HelpSteer2), revision `'+info['revision']+'`. Preference supervision uses independent human preference strength, never attribute-derived pair labels.\n\n'
    report+=f'Raw attribute rows: train {len(rawtrain)}, validation {len(rawtest)}; raw preference rows {len(prefs)}. Distinct normalized prompts {ex.prompt_id.nunique()}, distinct response identities {len(ex)}. Canonicalized official train/validation overlap: {len(overlap)}. If any overlap exists, test precedence moves all matching responses to test. Duplicate response ratings are averaged. Self-pairs and duplicate ordered preference rows are excluded.\n\n'
    report+=stats.to_markdown(index=False)+'\n\nOfficial validation becomes TEST; DEV is a deterministic 10% hash partition of remaining prompts. Canonical prompt IDs use Unicode NFKC and collapsed whitespace. All exact and canonical prompt boundaries pass leakage checks.\n\n'
    report+='## Actual schemas\n\n```json\n'+json.dumps({'attribute':list(rawtrain[0]),'preference':list(prefs[0])},indent=2)+'\n```\n\n'
    report+='## Preference distributions\n\n'+pd.crosstab(pair.strength,pair.split).to_markdown()+'\n\nNegative strength prefers response 1, positive prefers response 2. Zero-strength ties are excluded from binary BT training/evaluation, but their responses remain in representation and attribute evaluation.\n\n'
    report+='## Response character lengths (token lengths added separately)\n\n'+chars.to_frame().to_markdown()+'\n\n## Attribute distributions\n\n'+attrs.describe().to_markdown()+'\n\n## Attribute Pearson correlations (descriptive only)\n\n'+attrs.corr().to_markdown()+'\n\n'
    report+='Human attributes are stored in a separate table. Primary preference-head fitting reads only responses and preferences. The full-supervision upper bound alone reads training attribute values. Helpfulness is descriptive only, never a primary semantic coordinate.\n'
    Path('reports/dataset_report.md').write_text(report); stats.to_csv(out/'split_counts.csv',index=False)
    Path('outputs/dataset_manifest.json').write_text(json.dumps(dict(info,official_normalized_overlap=len(overlap),schemas={'attribute':list(rawtrain[0]),'preference':list(prefs[0])}),indent=2))
    print(stats.to_string(index=False),flush=True)

if __name__=='__main__':main()