Bases: TextGenerator
Qwen2-VL 2B image-caption generator.
Load the pinned Qwen2-VL model and processor.
Source code in modosaic/text/generators/impl/qwen_2_2b.py
23
24
25
26
27
28
29
30
31
32
33 | def __init__(self) -> None:
"""Load the pinned Qwen2-VL model and processor."""
super().__init__()
self.model_specs = HFModelSpec(
model_id="Qwen/Qwen2-VL-2B-Instruct",
revision="895c3a49bc3fa70a340399125c650a463535e71c",
dtype=self.dtype,
)
self.model = QWEN22B._load_model(self.model_specs)
self.model.to(self.device).eval()
self.processor = QWEN22B._load_tokenizer(self.model_specs)
|
generate
Generate a caption for one image record.
Source code in modosaic/text/generators/impl/qwen_2_2b.py
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67 | @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:
messages = QWEN22B._get_messages(tmp_image_path.resolve().as_posix())
inputs = self._get_inference_inputs(messages)
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)[0].strip()
logger.debug(f"Generated text for record sample {record.sample_id} with {len(caption)} characters")
return caption
|