| import argparse,json |
| from pathlib import Path |
| import torch |
| from tokenizers import Tokenizer |
| from model import Config,CodeSprout |
| def load(root,device): |
| config=Config(**json.loads((root/'config.json').read_text()));model=CodeSprout(config) |
| model.load_state_dict(torch.load(root/'model.pt',map_location='cpu',weights_only=True));model=model.to(device).eval();tok=Tokenizer.from_file(str(root/'tokenizer.json'));return model,tok |
| def complete(model,tok,prompt,language,tokens=120,temperature=.7): |
| normalized=prompt.rstrip();trailing=prompt[len(normalized):] |
| ids=torch.tensor([tok.encode('<|'+language+'|>\n'+normalized).ids],device=next(model.parameters()).device) |
| result=model.generate(ids,steps=tokens,temperature=temperature,eos=tok.token_to_id('<|endoftext|>')) |
| text=tok.decode(result[0,ids.shape[1]:].tolist()) |
| return text[len(trailing):] if trailing and text.startswith(trailing) else text |
| if __name__=='__main__': |
| p=argparse.ArgumentParser();p.add_argument('--checkpoint',type=Path,required=True);p.add_argument('--language',choices=['html','javascript','python','typescript'],default='python');p.add_argument('--prompt',default='def add(a, b):\n');p.add_argument('--tokens',type=int,default=120);a=p.parse_args();device='cuda' if torch.cuda.is_available() else 'cpu';model,tok=load(a.checkpoint,device);print(complete(model,tok,a.prompt,a.language,a.tokens)) |
|
|
