#!/usr/bin/env python3
"""Exercise the four-prompt T4 inference path on training prompts, never blind data."""
from pathlib import Path
import json
import sys
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / 'src'))
from unsloth import FastLanguageModel
import torch
from bitsandbytes.nn import Linear4bit
from generate import UnslothGenerator
from vibegame import BASE_MODEL

if not torch.cuda.is_available() or 'T4' not in torch.cuda.get_device_name():
    raise RuntimeError('This capacity check requires the already allocated free T4')
model, tokenizer = FastLanguageModel.from_pretrained(model_name=BASE_MODEL,
    revision='4858886f896bd05db823557fc5cbc4ac28342af7', max_seq_length=4096,
    dtype=torch.float16, load_in_4bit=True, trust_remote_code=False, use_exact_model_name=True)
for module in model.modules():
    if isinstance(module, Linear4bit):
        module.compute_dtype = torch.float16
        module.compute_type_is_set = True
FastLanguageModel.for_inference(model)
backend = UnslothGenerator.__new__(UnslothGenerator)
backend.model, backend.tokenizer, backend.torch = model, tokenizer, torch
prompts = [json.loads(line)['instruction'] for line in (ROOT/'data/teacher_prompts.jsonl').read_text().splitlines()][:4]
torch.cuda.reset_peak_memory_stats()
results = backend.generate_many(prompts, 3072)
report = {'kind':'untrained_base_capacity_smoke_on_training_prompts_not_model_eval',
          'batch_size':4, 'max_new_tokens':3072, 'peak_allocated_bytes':torch.cuda.max_memory_allocated(),
          'results': [{k:v for k,v in result.items() if k != 'output'} for result in results]}
target = Path(sys.argv[1])
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(json.dumps(report, indent=2)+'\n')
print(json.dumps(report, indent=2), flush=True)