| """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)) |
|
|
