shared_material/extract_shared_head.py

"""Append original BF16 no/yes rows to the frozen 86-row choice head; no training."""
import argparse
import hashlib
import json
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import load_file, save_file
from transformers import AutoTokenizer

p=argparse.ArgumentParser()
p.add_argument('--model',type=Path,default=Path('weights'))
p.add_argument('--choice-head',type=Path,default=Path('head/head.safetensors'))
p.add_argument('--out',type=Path,default=Path('shared_head'))
a=p.parse_args()
tokenizer=AutoTokenizer.from_pretrained(a.model,local_files_only=True)
encoded=[tokenizer.encode(word,add_special_tokens=False) for word in ['no','yes']]
if any(len(ids)!=1 for ids in encoded):raise ValueError('no/yes must be single tokens')
ids=torch.tensor([x[0] for x in encoded]);old=load_file(a.choice_head)
index=json.loads((a.model/'model.safetensors.index.json').read_text())['weight_map']
with safe_open(a.model/index['lm_head.weight'],framework='pt',device='cpu') as f:
 source=f.get_slice('lm_head.weight');rows=torch.cat([source[i:i+1] for i in ids.tolist()])
if rows.dtype!=torch.bfloat16:raise ValueError('This extractor requires original BF16 lm_head weights')
head={'weight':torch.cat([old['weight'],rows]),'token_ids':torch.cat([old['token_ids'],ids])}
if head['token_ids'].unique().numel()!=88:raise ValueError('Expected 86 distinct labels plus no/yes')
assert torch.equal(head['weight'][:86],old['weight'])
a.out.mkdir(parents=True,exist_ok=True)
save_file(head,a.out/'head.safetensors')
metadata={'tokens':dict(zip(['no','yes'],ids.tolist())),'shape':list(head['weight'].shape),'dtype':str(rows.dtype),
          'training':False,'original86rows_exact':True,'sha256':hashlib.sha256((a.out/'head.safetensors').read_bytes()).hexdigest()}
(a.out/'metadata.json').write_text(json.dumps(metadata,indent=2))
print(json.dumps(metadata,indent=2))