| | import os |
| | from typing import List |
| |
|
| | import torch |
| | import torch.nn as nn |
| | from einops import rearrange |
| | from transformers import AutoTokenizer, Qwen3ForCausalLM |
| |
|
| | from mmgp import offload |
| | from shared.utils import files_locator as fl |
| |
|
| | OUTPUT_LAYERS = [9, 18, 27] |
| | MAX_LENGTH = 512 |
| |
|
| |
|
| | class Qwen3Embedder(nn.Module): |
| | def __init__( |
| | self, |
| | model_spec: str | None = None, |
| | tokenizer_path: str | None = None, |
| | torch_dtype: str = "bfloat16", |
| | ): |
| | super().__init__() |
| | file_path = model_spec |
| | default_config = os.path.join(os.path.dirname(file_path), "config.json") |
| | self.model = offload.fast_load_transformers_model( |
| | file_path, |
| | writable_tensors=False, |
| | modelClass=Qwen3ForCausalLM, |
| | defaultConfigPath=default_config, |
| | ) |
| |
|
| | tokenizer_root = tokenizer_path or os.path.dirname(file_path) |
| | if tokenizer_root and not os.path.isabs(tokenizer_root): |
| | tokenizer_root = fl.locate_folder(tokenizer_root) |
| | tokenizer_subdir = os.path.join(tokenizer_root, "tokenizer") |
| | if os.path.isdir(tokenizer_subdir): |
| | tokenizer_root = tokenizer_subdir |
| | self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_root, trust_remote_code=True) |
| | self.max_length = MAX_LENGTH |
| |
|
| | @torch.no_grad() |
| | def forward(self, txt: List[str]): |
| | all_input_ids = [] |
| | all_attention_masks = [] |
| |
|
| | for prompt in txt: |
| | messages = [{"role": "user", "content": prompt}] |
| | text = self.tokenizer.apply_chat_template( |
| | messages, |
| | tokenize=False, |
| | add_generation_prompt=True, |
| | enable_thinking=False, |
| | ) |
| |
|
| | model_inputs = self.tokenizer( |
| | text, |
| | return_tensors="pt", |
| | padding="max_length", |
| | truncation=True, |
| | max_length=self.max_length, |
| | ) |
| | all_input_ids.append(model_inputs["input_ids"]) |
| | all_attention_masks.append(model_inputs["attention_mask"]) |
| |
|
| | input_ids = torch.cat(all_input_ids, dim=0).to(self.model.device) |
| | attention_mask = torch.cat(all_attention_masks, dim=0).to(self.model.device) |
| |
|
| | output = self.model( |
| | input_ids=input_ids, |
| | attention_mask=attention_mask, |
| | output_hidden_states=True, |
| | use_cache=False, |
| | ) |
| |
|
| | out = torch.stack([output.hidden_states[k] for k in OUTPUT_LAYERS], dim=1) |
| | return rearrange(out, "b c l d -> b l (c d)") |
| |
|
