Spaces:
Sleeping
Sleeping
File size: 3,811 Bytes
da3fe02 |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import Iterable, Optional
import torch
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
logger = logging.getLogger("labeler")
@dataclass
class ModelConfig:
model_id: str
device: str
precision: str
max_new_tokens: int
load_4bit: bool
class LabelerModel:
def __init__(self, config: ModelConfig) -> None:
self.config = config
self.device = _resolve_device(config.device)
self.dtype = _resolve_dtype(config.precision, self.device)
quantization_config = None
load_kwargs: dict[str, object] = {}
if config.load_4bit:
try:
from transformers import BitsAndBytesConfig
except ImportError as exc:
raise RuntimeError("bitsandbytes is required for 4-bit loading") from exc
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=self.dtype,
)
load_kwargs["quantization_config"] = quantization_config
load_kwargs["device_map"] = "auto"
elif self.device.startswith("cuda"):
load_kwargs["device_map"] = "auto"
self.processor = AutoProcessor.from_pretrained(config.model_id)
self.model = Qwen2VLForConditionalGeneration.from_pretrained(
config.model_id,
torch_dtype=self.dtype,
low_cpu_mem_usage=True,
**load_kwargs,
)
if not load_kwargs.get("device_map"):
self.model.to(self.device)
self.model.eval()
def generate_texts(
self,
messages_list: list[list[dict[str, object]]],
images: Optional[list[object]],
) -> list[str]:
prompts = [
self.processor.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
for messages in messages_list
]
if images is None:
inputs = self.processor(
text=prompts,
padding=True,
return_tensors="pt",
)
else:
inputs = self.processor(
text=prompts,
images=images,
padding=True,
return_tensors="pt",
)
inputs = _move_to_device(inputs, self.model.device)
with torch.inference_mode():
output_ids = self.model.generate(
**inputs,
max_new_tokens=self.config.max_new_tokens,
do_sample=False,
)
prompt_length = inputs["input_ids"].shape[1]
generated_ids = output_ids[:, prompt_length:]
return self.processor.batch_decode(generated_ids, skip_special_tokens=True)
def _resolve_device(device: str) -> str:
if device == "auto":
return "cuda" if torch.cuda.is_available() else "cpu"
return device
def _resolve_dtype(precision: str, device: str) -> torch.dtype:
if precision == "fp32":
return torch.float32
if precision == "bf16":
if device.startswith("cuda") and torch.cuda.is_bf16_supported():
return torch.bfloat16
return torch.float16
if precision == "fp16":
return torch.float16
if device.startswith("cuda"):
return torch.float16
return torch.float32
def _move_to_device(inputs: dict[str, object], device: torch.device | str) -> dict[str, object]:
moved = {}
for key, value in inputs.items():
if hasattr(value, "to"):
moved[key] = value.to(device)
else:
moved[key] = value
return moved
|