Bases: TextGenerator
Qwen2.5-VL 3B image-caption generator.
Load the pinned Qwen2.5-VL model and processor.
Source code in modosaic/text/generators/impl/qwen_2_5_3b.py
23
24
25
26
27
28
29
30
31
32
33 | def __init__(self) -> None:
"""Load the pinned Qwen2.5-VL model and processor."""
super().__init__()
self.model_specs = HFModelSpec(
model_id="Qwen/Qwen2.5-VL-3B-Instruct",
revision="66285546d2b821cf421d4f5eb2576359d3770cd3",
dtype=self.dtype,
)
self.model = QWEN253B._load_model(self.model_specs)
self.model.to(self.device).eval()
self.processor = QWEN253B._load_tokenizer(self.model_specs)
|
generate
Generate a caption for one image record.
Source code in modosaic/text/generators/impl/qwen_2_5_3b.py
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74 | @override
@torch.inference_mode
def generate(self, record: ImageRecord) -> str:
"""Generate a caption for one image record."""
logger.debug(f"Generating text for record sample {record.sample_id}")
with ImageService.temporary_image_file(record.image_bytes) as tmp_image_path:
file_uri = QWEN253B._get_file_uri(tmp_image_path)
messages = QWEN253B._get_messages(file_uri)
inference_text = self._get_inference_text(messages)
image_inputs, video_inputs = process_vision_info(messages)
inputs = self._get_inference_inputs(inference_text, image_inputs, video_inputs)
inputs = {k: v.to(self.device) for k, v in inputs.items()}
outputs = self.model.generate(**inputs, max_new_tokens=MAX_CAPTIONING_TOKENS)
gen = outputs[:, inputs["input_ids"].shape[1]:]
caption = self.processor.batch_decode(
gen,
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)[0].strip()
logger.debug(f"Generated text for record sample {record.sample_id} with {len(caption)} characters")
return caption
|