"""Simple causal completion CLI (no KV cache)."""
import argparse
from contextlib import nullcontext
import torch
from load_model import load_model

def main():
    p=argparse.ArgumentParser()
    p.add_argument('--model-dir',default='.')
    p.add_argument('--prompt',default='The scientific method is')
    p.add_argument('--max-new-tokens',type=int,default=64)
    p.add_argument('--temperature',type=float,default=0.0)
    p.add_argument('--device',default='cuda' if torch.cuda.is_available() else 'cpu')
    args=p.parse_args();torch.set_num_threads(4)
    model,tok=load_model(args.model_dir,args.device)
    ids=tok.encode(args.prompt).ids
    if not ids:raise ValueError('Prompt must encode to at least one token')
    eos=tok.token_to_id('<|endoftext|>')
    with torch.inference_mode():
        for _ in range(args.max_new_tokens):
            x=torch.tensor([ids[-model.block:]],device=args.device)
            ctx=torch.autocast('cuda',dtype=torch.bfloat16) if args.device.startswith('cuda') else nullcontext()
            with ctx:logits=model(x)[0][0,-1].float()
            nxt=int(logits.argmax()) if args.temperature<=0 else int(torch.multinomial(torch.softmax(logits/args.temperature,dim=-1),1))
            ids.append(nxt)
            if nxt==eos:break
    print(tok.decode(ids,skip_special_tokens=True))
if __name__=='__main__':main()