File size: 5,058 Bytes
33163f4 | 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 | """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:
# HF checkpoint'ит только в train()-режиме; у Ettin все dropout = 0.0 (в конфиге), поэтому
# численно это ничего не меняет. Пересчёт слоёв в backward: обучаемые активации Ettin
# занимали ~22 ГБ/GPU при 128 x 416 токенах без компиляции (OOM 2026-09-25).
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))
|