| """Score a native System One JSON request using a SecJev checkpoint.""" |
| import argparse, json, pathlib, time, os, sys |
| sys.path.insert(0,str(pathlib.Path(os.environ.get('SECJEV_WORK', './secjev-work')).resolve()/'source/kev')) |
| import torch |
| from kev.api import SystemOneRequest, to_record, to_answers |
| from kev.checkpoint import Checkpoint, LoadOptions |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--model', required=True, help='Local checkpoint directory or owner/repository@revision') |
| parser.add_argument('--request', required=True, type=pathlib.Path) |
| parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu') |
| parser.add_argument('--base', help='Optional local copy of the pinned Qwen backbone') |
| args = parser.parse_args() |
| request = SystemOneRequest.model_validate(json.loads(args.request.read_text())) |
| checkpoint = Checkpoint(args.model) |
| if args.base: |
| checkpoint.meta.base = args.base |
| checkpoint.meta.base_revision = None |
| torch.set_num_threads(4) |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cudnn.allow_tf32 = False |
| torch.backends.cuda.enable_flash_sdp(False) |
| torch.backends.cuda.enable_mem_efficient_sdp(False) |
| tokenizer, model = checkpoint.load(args.device, LoadOptions(dtype=torch.float32, merge=False, attn='sdpa')) |
| record, metadata = to_record(request) |
| encoded = model.encode(tokenizer, record, max_state=8192, max_branch=8192, strict=True) |
| from kev.model import rows_of |
| state, _, branches = rows_of(encoded) |
| if any(len(state) + len(branch['ids']) > 8192 for branch in branches): |
| raise ValueError('Each state + question must fit within 8192 tokens; inputs are never truncated.') |
| if str(args.device).startswith('cuda'): |
| torch.cuda.synchronize() |
| started = time.perf_counter() |
| with torch.inference_mode(): |
| logits = model.forward(encoded) |
| probabilities = [torch.softmax(z.float(), -1).cpu().tolist() for z in logits] |
| if str(args.device).startswith('cuda'): |
| torch.cuda.synchronize() |
| result = { |
| 'model': args.model, |
| 'answers': to_answers(probabilities, metadata), |
| 'probabilities': {m['id']: dict(zip(m['keys'], p)) for m, p in zip(metadata, probabilities)}, |
| 'temperature': model.head.temperature, |
| 'forward_ms': (time.perf_counter() - started) * 1000, |
| } |
| print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
