"""ScienceQA forced-choice inference with Transformers + AutoRound."""
import argparse
import io
import json
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument('--index', type=int, default=0)
    args = parser.parse_args()
    import torch
    from datasets import load_dataset, Image as DatasetImage
    from PIL import Image
    from transformers import AutoConfig, AutoProcessor, AutoRoundConfig, Mistral3ForConditionalGeneration

    encoder = args.model / 'text_encoder'
    cfg = json.loads((encoder / 'quantization_config.json').read_text())
    cfg.update(backend='auto_round:torch_mxfp8', extra_config={},
               block_name_to_quantize='model.language_model.layers,model.vision_tower.transformer.layers,model.multi_modal_projector')
    config = AutoConfig.from_pretrained(encoder, local_files_only=True)
    config.quantization_config = AutoRoundConfig(**cfg).to_dict()
    model, loading = Mistral3ForConditionalGeneration.from_pretrained(
        encoder, config=config, dtype=torch.bfloat16, device_map='cuda:0',
        attn_implementation='sdpa', local_files_only=True, output_loading_info=True,
    )
    model.eval()
    assert not any(loading.get(key) for key in ('missing_keys', 'unexpected_keys', 'mismatched_keys'))
    quantized = {name for name, module in model.named_modules() if hasattr(module, 'weight_scale')}
    assert len(quantized) == 451
    assert all('model.multi_modal_projector.' + name in quantized
               for name in ('linear_1', 'linear_2', 'patch_merger.merging_layer'))
    processor = AutoProcessor.from_pretrained(args.model / 'tokenizer', local_files_only=True)
    dataset = load_dataset('changwangss/ScienceQA-IMG-Test-2017',
                           revision='d21962594f7ae2cd2819fe48e95903d8a3aef5a5', split='test')
    row = dataset.cast_column('image', DatasetImage(decode=False))[args.index]
    image = Image.open(io.BytesIO(row['image']['bytes'])).convert('RGB')
    letters = 'ABCDE'[:len(row['choices'])]
    question = row['hint'] + '\n' + row['question'] + '\n' + '\n'.join(
        f'{letter}. {choice}' for letter, choice in zip(letters, row['choices']))
    prompt = '<s>[INST][IMG]' + question + '\nAnswer with only the correct option letter.[/INST]'
    inputs = processor(images=image, text=prompt, return_tensors='pt').to('cuda:0')
    inputs = {key: value.to(torch.bfloat16) if value.is_floating_point() else value
              for key, value in inputs.items()}
    ids = [processor.tokenizer.encode(letter, add_special_tokens=False) for letter in letters]
    assert all(len(token_ids) == 1 for token_ids in ids)
    with torch.inference_mode():
        logits = model(**inputs, logits_to_keep=1, use_cache=False).logits[0, -1].float()
    assert torch.isfinite(logits).all().item()
    prediction = int(logits[[token_ids[0] for token_ids in ids]].argmax().item())
    print('Question:', row['question'])
    print('Prediction:', letters[prediction], 'Correct:', prediction == row['answer'])


if __name__ == '__main__':
    main()