| """Stream the student's raw JSON response from a checkpoint; optionally render an MCP request.""" |
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
| import torch |
| from tokenizers import Tokenizer,decoders |
| from tinyquery.data import serialize,make_tools |
| from tinyquery.evaluate import load_model,check_action,mcp_request,parse_action |
|
|
|
|
| def default_context(backend): |
| import random |
| tools,_,_,_=make_tools(backend,random.Random(4),'train','demo') |
| return {'backend':backend,'project_id':'demo', |
| 'schema':['CREATE TABLE customers (id INTEGER PRIMARY KEY, name TEXT, city TEXT, amount REAL);'], |
| 'tools':tools,'policy':'Read-only database access. Use provided tools. Ask when required information is missing.'} |
|
|
|
|
| def stream_response(model,tokenizer,prompt,max_tokens=160,stats=None): |
| ids=tokenizer.encode(prompt).ids |
| if len(ids)+max_tokens>model.config.context: raise ValueError('Prompt plus output budget exceeds model context') |
| device=next(model.parameters()).device; tokens=torch.tensor([ids],device=device) |
| past=None; decoder=decoders.DecodeStream(skip_special_tokens=True) |
| generated=[]; emitted='';start=time.perf_counter();first_token=None |
| with torch.inference_mode(): |
| for _ in range(max_tokens): |
| logits,past,_=model(tokens,past=past,use_cache=True,last_only=True) |
| token=int(logits[0,-1].argmax()); generated.append(token) |
| if first_token is None:first_token=time.perf_counter()-start |
| text=decoder.step(tokenizer,token) |
| if text: emitted+=text; yield text |
| if token==tokenizer.token_to_id('<|end|>'): break |
| tokens=torch.tensor([[token]],device=device) |
| full=tokenizer.decode(generated,skip_special_tokens=True) |
| if full.startswith(emitted) and len(full)>len(emitted): yield full[len(emitted):] |
| if stats is not None: |
| seconds=time.perf_counter()-start |
| stats.update(generated_tokens=len(generated),seconds=seconds,tokens_per_second=len(generated)/seconds, |
| first_token_seconds=first_token,device=str(device),includes_model_loading=False) |
|
|
|
|
| def main(): |
| p=argparse.ArgumentParser(); p.add_argument('question'); p.add_argument('--checkpoint',required=True) |
| p.add_argument('--tokenizer'); p.add_argument('--context',help='JSON file with backend, schema and MCP-style tools') |
| p.add_argument('--backend',choices=['mysql','supabase'],default='supabase') |
| p.add_argument('--tokens',type=int,default=160); p.add_argument('--mcp',action='store_true') |
| p.add_argument('--stats',action='store_true',help='Print measured generation speed to stderr after streaming') |
| args=p.parse_args(); device='cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu') |
| torch.set_num_threads(4) |
| tokenizer=Tokenizer.from_file(args.tokenizer or str(Path(args.checkpoint).parent/'tokenizer.json')) |
| model=load_model(args.checkpoint,device) |
| context=json.loads(Path(args.context).read_text()) if args.context else default_context(args.backend) |
| answer='';stats={} |
| for text in stream_response(model,tokenizer,serialize(context,args.question),args.tokens,stats): |
| print(text,end='',flush=True); answer+=text |
| print(flush=True) |
| if args.stats:print(json.dumps(stats),file=sys.stderr) |
| try: |
| action=check_action(parse_action(answer),context) |
| if args.mcp and action['action']=='call': print(json.dumps(mcp_request(action,context),ensure_ascii=False,indent=2)) |
| except Exception as exc: |
| print('Output validation failed:',str(exc),file=sys.stderr) |
| raise SystemExit(1) |
|
|
|
|
| if __name__=='__main__': main() |
|
|
