training/verify_observation_program_export.py

"""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()