| import os |
|
|
| import torch |
| from torch import Tensor, nn |
| from transformers import ( |
| CLIPTextModel, |
| CLIPTokenizer, |
| T5EncoderModel, |
| T5Tokenizer, |
| __version__, |
| ) |
| from transformers.utils.quantization_config import QuantoConfig, BitsAndBytesConfig |
|
|
| CACHE_DIR = os.environ.get("HF_HOME", "~/.cache/huggingface") |
|
|
|
|
| def auto_quantization_config( |
| quantization_dtype: str, |
| ) -> QuantoConfig | BitsAndBytesConfig: |
| if quantization_dtype == "qfloat8": |
| return QuantoConfig(weights="float8") |
| elif quantization_dtype == "qint4": |
| return BitsAndBytesConfig( |
| load_in_4bit=True, |
| bnb_4bit_compute_dtype=torch.bfloat16, |
| bnb_4bit_quant_type="nf4", |
| ) |
| elif quantization_dtype == "qint8": |
| return BitsAndBytesConfig(load_in_8bit=True, llm_int8_has_fp16_weight=False) |
| elif quantization_dtype == "qint2": |
| return QuantoConfig(weights="int2") |
| elif quantization_dtype is None or quantization_dtype == "bfloat16": |
| return None |
| else: |
| raise ValueError(f"Unsupported quantization dtype: {quantization_dtype}") |
|
|
|
|
| class HFEmbedder(nn.Module): |
| def __init__( |
| self, |
| version: str, |
| max_length: int, |
| device: torch.device | int, |
| quantization_dtype: str | None = None, |
| offloading_device: torch.device | int | None = torch.device("cpu"), |
| is_clip: bool = False, |
| **hf_kwargs, |
| ): |
| super().__init__() |
| self.offloading_device = ( |
| offloading_device |
| if isinstance(offloading_device, torch.device) |
| else torch.device(offloading_device) |
| ) |
| self.device = ( |
| device if isinstance(device, torch.device) else torch.device(device) |
| ) |
| self.is_clip = version.startswith("openai") or is_clip |
| self.max_length = max_length |
| self.output_key = "pooler_output" if self.is_clip else "last_hidden_state" |
|
|
| auto_quant_config = ( |
| auto_quantization_config(quantization_dtype) |
| if quantization_dtype is not None |
| and quantization_dtype != "bfloat16" |
| and quantization_dtype != "float16" |
| else None |
| ) |
|
|
| |
| if isinstance(auto_quant_config, BitsAndBytesConfig): |
| hf_kwargs["device_map"] = {"": self.device.index} |
| if auto_quant_config is not None: |
| hf_kwargs["quantization_config"] = auto_quant_config |
|
|
| if self.is_clip: |
| self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained( |
| version, max_length=max_length |
| ) |
|
|
| self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained( |
| version, |
| **hf_kwargs, |
| ) |
|
|
| else: |
| self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained( |
| version, max_length=max_length |
| ) |
| self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained( |
| version, |
| **hf_kwargs, |
| ) |
|
|
| def offload(self): |
| self.hf_module.to(device=self.offloading_device) |
| torch.cuda.empty_cache() |
|
|
| def cuda(self): |
| self.hf_module.to(device=self.device) |
|
|
| def forward(self, text: list[str]) -> Tensor: |
| batch_encoding = self.tokenizer( |
| text, |
| truncation=True, |
| max_length=self.max_length, |
| return_length=False, |
| return_overflowing_tokens=False, |
| padding="max_length", |
| return_tensors="pt", |
| ) |
| outputs = self.hf_module( |
| input_ids=batch_encoding["input_ids"].to(self.hf_module.device), |
| attention_mask=None, |
| output_hidden_states=False, |
| ) |
| return outputs[self.output_key] |
|
|
|
|
| if __name__ == "__main__": |
| model = HFEmbedder( |
| "city96/t5-v1_1-xxl-encoder-bf16", |
| max_length=512, |
| device=0, |
| quantization_dtype="qfloat8", |
| ) |
| o = model(["hello"]) |
| print(o) |
|
|
