File size: 6,916 Bytes
b6d3dd9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """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()
|