predict_v2.py

"""Compile one quote into an unapproved, schema-grounded proposal. Never executes it."""
import argparse
import json
from pathlib import Path
import re
from runtime_contract import validate_input, parse_proposal, strict_json
from train_v2 import SYSTEM_V2, REVISION

ADAPTER_REVISION = 'a897deec75477ffb27e4a1ef9df8d4b52fb6bdf8'

def predict(quote,schema,adapter='khursheed/datacard-ci',adapter_revision=None,
            device='cpu',local_files_only=False):
    validate_input(quote, schema)
    if device not in ('cpu', 'cuda'):
        raise ValueError('Choose cpu or cuda explicitly')
    if not Path(adapter).is_dir():
        adapter_revision = adapter_revision or (ADAPTER_REVISION if adapter == 'khursheed/datacard-ci' else None)
        if not isinstance(adapter_revision, str) or not re.fullmatch(r'[0-9a-f]{40}', adapter_revision):
            raise ValueError('Remote adapters require a pinned 40-character commit')
    import torch
    from transformers import AutoTokenizer,AutoModelForCausalLM
    from peft import PeftModel
    if device == 'cuda' and not torch.cuda.is_available():
        raise ValueError('CUDA requested but unavailable')
    tokenizer=AutoTokenizer.from_pretrained('Qwen/Qwen3-0.6B',revision=REVISION,trust_remote_code=False,local_files_only=local_files_only)
    messages=[{'role':'system','content':SYSTEM_V2},{'role':'user','content':json.dumps({'quote':quote,'schema':schema},separators=(',',':'))}]
    text=tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True,enable_thinking=False)
    inputs=tokenizer(text,return_tensors='pt',add_special_tokens=False).to(device)
    if inputs.input_ids.shape[1]+128>640:
        raise ValueError('Input exceeds prototype context budget; select a shorter quote')
    # Reject oversized prompts before allocating model weights.
    base=AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-0.6B',revision=REVISION,trust_remote_code=False,dtype=torch.float32,local_files_only=local_files_only).to(device)
    kwargs={'revision':adapter_revision} if adapter_revision else {}
    model=PeftModel.from_pretrained(base,adapter,local_files_only=local_files_only,**kwargs).eval()
    with torch.no_grad():
        if device=='cuda':
            with torch.autocast('cuda',dtype=torch.float16):
                output=model.generate(**inputs,max_new_tokens=128,do_sample=False,pad_token_id=tokenizer.eos_token_id)
        else:
            output=model.generate(**inputs,max_new_tokens=128,do_sample=False,pad_token_id=tokenizer.eos_token_id)
    raw=tokenizer.decode(output[0,inputs.input_ids.shape[1]:],skip_special_tokens=True).strip()
    return parse_proposal(raw, quote, schema)

if __name__=='__main__':
    p=argparse.ArgumentParser()
    p.add_argument('--quote',required=True)
    p.add_argument('--schema',required=True,help='JSON mapping split names to column lists')
    p.add_argument('--adapter',default='khursheed/datacard-ci')
    p.add_argument('--adapter-revision')
    p.add_argument('--device',choices=['cpu','cuda'],default='cpu')
    p.add_argument('--offline',action='store_true',help='Use only locally cached model files')
    a=p.parse_args()
    print(json.dumps(predict(a.quote,strict_json(a.schema),a.adapter,a.adapter_revision,device=a.device,local_files_only=a.offline),indent=2))