File size: 2,871 Bytes
b296ad4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | """Export only inference weights and public project code; never upload credentials or remote config."""
import argparse
import hashlib
import json
from pathlib import Path
import shutil
import torch
from safetensors.torch import save_file
from tinyquery.model import Config,TinyQuery
def main():
p=argparse.ArgumentParser(); p.add_argument('--checkpoint',required=True); p.add_argument('--data',required=True)
p.add_argument('--out',required=True); args=p.parse_args()
source=Path(args.checkpoint); data=Path(args.data); out=Path(args.out); out.mkdir(parents=True,exist_ok=True)
if source.suffix=='.safetensors':
config=json.loads((source.parent/'config.json').read_text())
shutil.copy2(source,out/'model.safetensors')
info=source.parent/'checkpoint-info.json'
if not info.exists():info=source.parent/'best-info.json'
if info.exists(): shutil.copy2(info,out/'checkpoint-info.json')
else:
checkpoint=torch.load(source,map_location='cpu',weights_only=False)
config=checkpoint['config']
state={k:v.to(torch.bfloat16).contiguous() for k,v in checkpoint['model'].items()}
save_file(state,str(out/'model.safetensors'))
info={k:checkpoint[k] for k in ['step','processed_tokens','response_tokens','training_seconds','random_initialization']}
(out/'checkpoint-info.json').write_text(json.dumps(info,indent=2))
(out/'config.json').write_text(json.dumps(config,indent=2))
shutil.copy2(data/'tokenizer.json',out/'tokenizer.json')
package=Path(__file__).parent
shutil.copytree(package,out/'tinyquery',dirs_exist_ok=True,ignore=shutil.ignore_patterns('__pycache__','*.pyc'))
shutil.copy2(package/'requirements-inference.txt',out/'requirements.txt')
for name in ['tokenization.json','grounding-stats.json','split-audit.json']:
if (data/name).exists(): shutil.copy2(data/name,out/name)
model=TinyQuery(Config(**config)); count=sum(p.numel() for p in model.parameters())
from safetensors.torch import load_file
model.load_state_dict(load_file(str(out/'model.safetensors')),strict=True)
manifest={'parameters':count,'format':'Custom native PyTorch; see tinyquery/model.py, not a Transformers AutoModel checkpoint',
'random_initialization':True,'files':{}}
for file in sorted(out.rglob('*')):
if file.is_file() and file.name!='manifest.json':
digest=hashlib.sha256()
with file.open('rb') as stream:
for block in iter(lambda:stream.read(4*1024*1024),b''): digest.update(block)
manifest['files'][str(file.relative_to(out))]={'bytes':file.stat().st_size,'sha256':digest.hexdigest()}
(out/'manifest.json').write_text(json.dumps(manifest,indent=2)); print(json.dumps({'parameters':count,'files':len(manifest['files'])}))
if __name__=='__main__': main()
|