| """Ettin-68M (jhu-clsp/ettin-encoder-68m, MIT) — ModernBERT-энкодер, hidden 512, byte-level BPE. |
| |
| Взят из релиза Agate. Из файла релиза оставлен только Ettin: вариантов T5 и Hash-заглушки из |
| FCDM/T2I-ветки (и функции для разморозки top-N блоков T5) в проекте нет — модель заморожена. |
| |
| Держит регистр, пунктуацию, диакритику и нелатинские названия, которые SentencePiece у Flan-T5 |
| превращает в <unk>. Пустая строка кодируется как [CLS][SEP] — это безусловный контекст для CFG. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
|
|
|
|
| def _load_modernbert(path: str): |
| """AutoModel.from_pretrained, или — там, где transformers отказывается читать .bin на torch < 2.6 |
| (Ettin поставляется только с pytorch_model.bin) — сборка из конфига и загрузка state_dict с |
| weights_only=True. Тензоры MLM-головы отбрасываются.""" |
| from transformers import AutoConfig, AutoModel |
| try: |
| return AutoModel.from_pretrained(path, torch_dtype=torch.float32) |
| except ValueError: |
| import os |
| m = AutoModel.from_config(AutoConfig.from_pretrained(path)) |
| f = os.path.join(path, "pytorch_model.bin") if os.path.isdir(path) else None |
| if f is None: |
| from huggingface_hub import hf_hub_download |
| f = hf_hub_download(path, "pytorch_model.bin") |
| sd = {k.removeprefix("model."): v for k, v in torch.load(f, map_location="cpu", weights_only=True).items()} |
| missing, _ = m.load_state_dict(sd, strict=False) |
| if missing: |
| raise KeyError(f"{path}: encoder tensors missing from the checkpoint: {missing[:5]}") |
| return m |
|
|
|
|
| class EttinTextEncoder: |
| """ModernBERT-энкодер Agate. train_blocks != 0 тренирует весь энкодер (опция «верхние блоки» |
| бессмысленна: он маленький, и пространство обусловливания всё равно учится заново). |
| |
| Масок паддинга две: attention_mask в модели и mask наружу (1 = настоящий токен), по нему |
| кросс-внимание мышления игнорирует паддинг.""" |
|
|
| def __init__(self, path: str, device, max_len: int = 128, train_blocks: int = 0, grad_ckpt: bool = False): |
| from transformers import AutoTokenizer |
| self.device, self.max_len = torch.device(device), max_len |
| self.tok = AutoTokenizer.from_pretrained(path) |
| self.model = _load_modernbert(path).to(self.device) |
| self.model.eval().requires_grad_(bool(train_blocks)) |
| if grad_ckpt and train_blocks: |
| |
| |
| |
| if any(getattr(self.model.config, k, 0) for k in ("attention_dropout", "mlp_dropout", "embedding_dropout")): |
| raise ValueError("Ettin checkpointing needs train() mode, but this config has dropout") |
| self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) |
| self.model.train() |
| self.dim = self.model.config.hidden_size |
| self.trainable = [(n, p) for n, p in self.model.named_parameters() if p.requires_grad] |
|
|
| def tokenize(self, captions: list[str], pad_to: int = 1) -> tuple[torch.Tensor, torch.Tensor]: |
| t = self.tok(captions, padding="longest", truncation=True, max_length=self.max_len, |
| return_tensors="pt") |
| ids, mask = t["input_ids"], t["attention_mask"] |
| extra = (-ids.shape[1]) % pad_to |
| if extra: |
| ids = torch.nn.functional.pad(ids, (0, extra), value=self.tok.pad_token_id) |
| mask = torch.nn.functional.pad(mask, (0, extra)) |
| return ids, mask |
|
|
| def encode(self, ids: torch.Tensor, mask: torch.Tensor, grad: bool = False |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| ids, mask = ids.to(self.device, non_blocking=True), mask.to(self.device, non_blocking=True) |
| with torch.set_grad_enabled(grad and bool(self.trainable)), \ |
| torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda"): |
| h = self.model(input_ids=ids, attention_mask=mask).last_hidden_state |
| return h.float(), mask |
|
|
| def __call__(self, captions: list[str]) -> tuple[torch.Tensor, torch.Tensor]: |
| return self.encode(*self.tokenize(captions)) |
|
|
