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