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