| from pathlib import Path |
| import argparse,json,hashlib |
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
| p=argparse.ArgumentParser();p.add_argument('--root',required=True);p.add_argument('--revision',required=True);a=p.parse_args();root=Path(a.root) |
| out=root/'mlx-decoder';out.mkdir(exist_ok=True) |
| weights={} |
| for path in sorted((root/'upstream').glob('*.safetensors')): |
| if path.name=='feature_extractor.safetensors':continue |
| with safe_open(str(path),framework='pt',device='cpu') as f: |
| for key in f.keys(): |
| if key.startswith('model.decoder.') or key.startswith('lm_head.'): |
| weights[key.removeprefix('model.decoder.')]=f.get_tensor(key).half().contiguous() |
| if 'lm_head.weight' not in weights:weights['lm_head.weight']=weights['embedding.token_embedding.weight'].clone() |
| assert len(weights)>500, len(weights) |
| save_file(weights,str(out/'decoder.safetensors')) |
| cfg=json.loads((root/'upstream/config.json').read_text()) |
| minimal={k:cfg[k] for k in ['d_model','decoder_layers','decoder_attention_heads','decoder_ffn_dim','vocab_size','max_target_positions','tie_word_embeddings']} |
| minimal.update(model_type='bodhan_mlx_decoder',architectures=['BodhanMLXDecoder'],layer_norm_eps=1e-5) |
| (out/'config.json').write_text(json.dumps(minimal,indent=2)+'\n') |
| h=hashlib.sha256() |
| with (out/'decoder.safetensors').open('rb') as f: |
| for chunk in iter(lambda:f.read(8*1024*1024),b''):h.update(chunk) |
| r={'upstream_model':'bodhan-ai/indic-transcribe-'+('flex' if 'flex' in root.name else 'core'),'revision':a.revision,'dtype':'float16','tensors':len(weights),'bytes':(out/'decoder.safetensors').stat().st_size,'sha256':h.hexdigest(),'keys':list(weights)} |
| (out/'manifest.json').write_text(json.dumps(r,indent=2));print({k:v for k,v in r.items() if k!='keys'},flush=True) |
|
|
