"""Reload an exported Edge VLM and preserve one raw generation for each Tancho role.""" import argparse import gc import hashlib import json from pathlib import Path import sys import time ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from training.observation_program_loader import ObservationProgramDataset, verify_pack def _tensor_hash(tensor): tensor = tensor.detach().contiguous().cpu() digest = hashlib.sha256(str((tuple(tensor.shape), str(tensor.dtype))).encode()) digest.update(tensor.view(-1).view(__import__('torch').uint8).numpy().tobytes()) return digest.hexdigest() def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--pack', type=Path, required=True) parser.add_argument('--base', type=Path, required=True) parser.add_argument('--export', type=Path, required=True) parser.add_argument('--output', type=Path, required=True) parser.add_argument('--purpose', choices=('training', 'compatibility'), required=True) parser.add_argument('--max-new-tokens', type=int, default=512) args = parser.parse_args() if args.output.exists(): raise ValueError('Use a fresh export verification directory') manifest = verify_pack(args.pack, purpose=args.purpose) args.output.mkdir(parents=True) import torch from transformers import AutoModelForImageTextToText import cosmos_framework.model.generator.reasoner.cosmos3_edge # registers native Auto class from cosmos_framework.data.generator.processors import build_processor from tancho.edge_chat import install from tancho.observation_intent import parse_intent_output from tancho.observation_program import parse_program_output if not torch.cuda.is_available() or torch.cuda.device_count() != 1: raise ValueError('Export smoke requires exactly one CUDA GPU') install(); torch.manual_seed(42) started = time.monotonic() model, loading = AutoModelForImageTextToText.from_pretrained( str(args.export), torch_dtype=torch.bfloat16, device_map='cuda:0', attn_implementation='sdpa', local_files_only=True, output_loading_info=True) if any(loading.get(key) for key in ('missing_keys', 'unexpected_keys', 'mismatched_keys', 'error_msgs')): raise ValueError('HF export reload changed model tensors: ' + repr(loading)) model.eval(); processor = build_processor(tokenizer_type=str(args.base), config_variant='hf') dataset = ObservationProgramDataset(args.pack, 'validation', purpose=args.purpose) selected = {} for index, row in enumerate(dataset.rows): selected.setdefault(row['role'], (index, row)) if set(selected) != {'reasoner', 'generator'}: raise ValueError('Validation split must contain both roles') records = [] for role in ('reasoner', 'generator'): index, row = selected[role]; item = dataset[index] if 'context' not in row: raise ValueError('Pack lacks validator context; use a context-bearing evaluation pack') inputs = processor.apply_chat_template([item['texts'][0]], tokenize=True, add_generation_prompt=True, return_tensors='pt') tensor_inputs = {key: (value.unsqueeze(0) if key in ('input_ids', 'attention_mask') and value.ndim == 1 else value).to('cuda:0') for key, value in inputs.items() if torch.is_tensor(value)} torch.cuda.reset_peak_memory_stats(); generation_started = time.monotonic() with torch.inference_mode(): generated = model.generate(**tensor_inputs, max_new_tokens=args.max_new_tokens, do_sample=False, use_cache=False, eos_token_id=11, pad_token_id=0) torch.cuda.synchronize() tokens = generated[0, tensor_inputs['input_ids'].shape[-1]:].detach().cpu().tolist() if not tokens: raise ValueError(role + ' returned no tokens') raw = processor.processor.tokenizer.decode(tokens, skip_special_tokens=True) strict_valid = False; strict_error = None try: parsed = (parse_intent_output if role == 'reasoner' else parse_program_output)(raw) if role == 'reasoner': from tancho.observation_intent import validate_and_bind_intents validate_and_bind_intents(parsed, row['context'], now_ms=row['context']['captured_at_ms'], raw_text=raw) else: from tancho.observation_program import validate_observation_program validate_observation_program(parsed, row['context'], row['verified_intents'], now_ms=row['context']['captured_at_ms'], raw_text=raw) strict_valid = True except Exception as exc: strict_error = type(exc).__name__ + ': ' + str(exc) raw_path = args.output / (role + '.txt'); raw_path.write_text(raw) records.append({'role': role, 'sample_id': row['sample_id'], 'generation_completed': True, 'single_attempt': True, 'strict_parse_valid': strict_valid, 'strict_parse_error': strict_error, 'raw_path': raw_path.name, 'raw_sha256': hashlib.sha256(raw.encode()).hexdigest(), 'token_ids_sha256': hashlib.sha256(json.dumps(tokens).encode()).hexdigest(), 'returned_tokens': len(tokens), 'input_tensor_sha256': {key: _tensor_hash(value) for key, value in tensor_inputs.items()}, 'generation_seconds': time.monotonic() - generation_started, 'peak_allocated_bytes': torch.cuda.max_memory_allocated()}) del tensor_inputs, generated report = {'schema_version': 'tancho-observation-program-export-smoke-1.0', 'passed': True, 'pack_manifest_sha256': manifest['manifest_sha256'], 'model_revision': manifest['model_revision'], 'hf_export_reloaded': True, 'loading_info': loading, 'single_attempt': True, 'raw_outputs_preserved': True, 'reasoner_generation_completed': True, 'generator_generation_completed': True, 'records': records, 'elapsed_seconds': time.monotonic() - started} (args.output / 'report.json').write_text(json.dumps(report, ensure_ascii=False, indent=2) + '\n') del model; gc.collect(); torch.cuda.empty_cache() print('TANCHO_OBSERVATION_PROGRAM_EXPORT ' + json.dumps({ 'passed': True, 'roles': [row['role'] for row in records], 'strict_parse_valid': {row['role']: row['strict_parse_valid'] for row in records}}, sort_keys=True)) if __name__ == '__main__': main()