File size: 3,158 Bytes
54daaf9 | 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 | """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() |