Spaces:
Running on Zero
Running on Zero
multimodalart HF Staff
GLiDE-0.4B decision model demo: vendored glide package, app, examples from model card
c1be530 verified | """Encoder backbone (ModernBERT architecture) for Glide checkpoints whose decision_config.json has "backbone": "encoder". | |
| The prompt TEXT is the same as for the decoder Glide models: it is rendered with the checkpoint's render tokenizer (render_tokenizer/: chat | |
| template and single-token option codes), then tokenised by the encoder tokenizer as [CLS] + ids + [SEP]. Prompts longer than "max_len" tokens keep | |
| the first "trunc_head" and the last (max_len - 2 - trunc_head) tokens. Every layer attends bidirectionally over the non-pad tokens; the readout reads | |
| the last real token before [SEP] (the final "(" of "Answer: (", pooling "last") or [CLS] (pooling "cls"). | |
| Tensor names (HF convention, no backbone.* keys): model.embeddings.*, model.layers.<i>.*, model.final_norm.* (as in ModernBertForMaskedLM; the MLM | |
| head is not included) in model-*.safetensors; the 255-way decision head in readout.safetensors {"weight": [255, hidden]}. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import torch | |
| ENC_PREFIX = "model." | |
| RENDER_DIR = "render_tokenizer" | |
| MAX_OPTIONS = 255 | |
| class EncoderTokenizer: | |
| """Prompt-side calls (chat template, vocab, encode of option codes) go to the render tokenizer; __call__ returns encoder ids.""" | |
| def __init__(self, render, enc, max_len, trunc_head): | |
| self.render, self.enc = render, enc | |
| self.max_len, self.trunc_head = int(max_len), int(trunc_head) | |
| self.cls_id, self.sep_id, self.pad_token_id = enc.cls_token_id, enc.sep_token_id, enc.pad_token_id | |
| def apply_chat_template(self, *args, **kwargs): | |
| return self.render.apply_chat_template(*args, **kwargs) | |
| def get_vocab(self): | |
| return self.render.get_vocab() | |
| def encode(self, text, add_special_tokens=False, **kwargs): | |
| return self.render.encode(text, add_special_tokens=add_special_tokens, **kwargs) | |
| def _wrap(self, ids): | |
| n = self.max_len - 2 | |
| if len(ids) > n: | |
| ids = ids[: self.trunc_head] + ids[len(ids) - (n - self.trunc_head):] | |
| return [self.cls_id] + list(ids) + [self.sep_id] | |
| def __call__(self, texts, add_special_tokens=False, **kwargs): | |
| ids = self.enc(list(texts), add_special_tokens=False)["input_ids"] | |
| return {"input_ids": [self._wrap(x) for x in ids]} | |
| class EncoderNetwork(torch.nn.Module): | |
| def __init__(self, config, pooling="last"): | |
| from transformers.models.modernbert.modeling_modernbert import ModernBertModel | |
| super().__init__() | |
| self.model = ModernBertModel(config) | |
| self.readout = torch.nn.Linear(config.hidden_size, MAX_OPTIONS, bias=False) | |
| self.pooling = pooling | |
| self.visual = None | |
| def forward(self, input_ids, padding, pixel_values=None, image_grid_thw=None): | |
| if pixel_values is not None: | |
| raise ValueError("encoder Glide checkpoints are text-only") | |
| pos = (padding.long().cumsum(-1) - 1).clamp_(min=0) | |
| h = self.model(input_ids=input_ids, attention_mask=padding, position_ids=pos).last_hidden_state | |
| if self.pooling == "cls": | |
| idx = (padding.shape[1] - padding.long().sum(-1)).clamp_(max=padding.shape[1] - 1) | |
| v = h[torch.arange(h.shape[0], device=h.device), idx] | |
| else: | |
| v = h[:, -2] # left padding: [..., "(", [SEP]] | |
| return self.readout(v) | |
| def load(path, config, device, dtype): | |
| """-> (EncoderNetwork in eval mode, EncoderTokenizer).""" | |
| from safetensors import safe_open | |
| from safetensors.torch import load_file | |
| from transformers import AutoConfig, AutoTokenizer | |
| render = AutoTokenizer.from_pretrained(os.path.join(path, config.get("render_tokenizer") or RENDER_DIR)) | |
| enc = AutoTokenizer.from_pretrained(path) | |
| tok = EncoderTokenizer(render, enc, config.get("max_len", 7999), config.get("trunc_head", 1024)) | |
| cfg = AutoConfig.from_pretrained(path) | |
| cfg._attn_implementation = "sdpa" | |
| net = EncoderNetwork(cfg, config.get("pooling", "last")).to(device=device, dtype=dtype) | |
| with open(os.path.join(path, "model.safetensors.index.json")) as f: | |
| files = sorted(set(json.load(f)["weight_map"].values())) | |
| where = {} | |
| handles = {f: safe_open(os.path.join(path, f), framework="pt", device="cpu") for f in files} | |
| for f, h in handles.items(): | |
| for k in h.keys(): | |
| where[k] = f | |
| missing, used = [], set() | |
| with torch.no_grad(): | |
| for n, p in net.model.named_parameters(): | |
| k = ENC_PREFIX + n | |
| if k not in where: | |
| missing.append(k) | |
| continue | |
| p.copy_(handles[where[k]].get_tensor(k).to(device=device, dtype=dtype)) | |
| used.add(k) | |
| net.readout.weight.copy_(load_file(os.path.join(path, "readout.safetensors"))["weight"].to(device=device, dtype=dtype)) | |
| if missing: | |
| raise RuntimeError(f"missing encoder weights: {missing[:5]} ({len(missing)} total)") | |
| extra = sorted(set(where) - used) | |
| if extra: | |
| raise RuntimeError(f"unused tensors in the checkpoint: {extra[:5]} ({len(extra)} total)") | |
| net.eval().requires_grad_(False) | |
| return net, tok | |