"""Standalone Hugging Face model for the Lumen alignment. This file is **copied verbatim into the released Hub repository** and is loaded there via ``trust_remote_code=True``. It therefore must not import ``lumen`` or anything else from this repository. Everything it needs is either in the released repo or pulled from the Hub at load time. The released repository carries only the alignment: rank-4 LoRA adapters on both encoders, two projection heads and a learned logit scale, 2,984,961 parameters in total. The Virchow2 and BioMedBERT backbones are **not** redistributed here. They are fetched from their own repositories when the model is instantiated, so their licences and access controls apply directly to the user. Usage:: from transformers import AutoModel model = AutoModel.from_pretrained("digitalpathologybern/Lumen", trust_remote_code=True).eval() tokens = model.tokenizer(["metastatic carcinoma in a lymph node"], padding=True, truncation=True, max_length=128, return_tensors="pt") out = model(images=pixels, tokens=tokens) # pixels: (B, 3, 224, 224) Loading prints a long "newly initialized" warning naming every backbone tensor. That is expected and is not a problem: the released repository withholds the backbones on purpose, and transformers reports the resulting missing keys. The weights are the pretrained ones, fetched in ``__init__`` before ``from_pretrained`` looks for them, and nothing re-initialises them. """ from __future__ import annotations import torch import torch.nn as nn from transformers import PretrainedConfig, PreTrainedModel VISION_REPO = "paige-ai/Virchow2" TEXT_REPO = "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext" # Virchow2 is a gated repository: access must be requested and granted on the # Hub before the weights can be downloaded. A bare HTTP 401/403 from deep inside # timm is not something a user can act on, so every backbone fetch is wrapped and # re-raised with the repository URL and what to do about it. _ACCESS_HELP = """ Could not download the {name} backbone from '{repo}'. Lumen is an adapter. It carries only the 2,984,961-parameter alignment and does not redistribute either backbone, so both are fetched from their own repositories at load time. {reason} To fix this: 1. Open https://huggingface.co/{repo} and, if the page shows a licence or access form, accept it and wait for access to be granted. 2. Authenticate this machine: huggingface-cli login (or set the HF_TOKEN environment variable to a token with read access) 3. Retry. If you already have the weights locally, pass the path directly: AutoModel.from_pretrained(..., trust_remote_code=True, {kwarg}="/path/to/{leaf}") """ def _diagnose(exc: Exception) -> str: """Say why the download failed, in terms the caller can act on. Classified by exception type first. The hub's message text mentions gating even for a plain 404, because the same help string covers private and gated repositories, so matching on the word "gated" misreports a typo as an access problem. """ try: from huggingface_hub.utils import (EntryNotFoundError, GatedRepoError, RepositoryNotFoundError) except ImportError: # pragma: no cover - very old hub GatedRepoError = RepositoryNotFoundError = EntryNotFoundError = () if GatedRepoError and isinstance(exc, GatedRepoError): return ("The repository is gated and this account has not been granted " "access yet.") if RepositoryNotFoundError and isinstance(exc, RepositoryNotFoundError): return ("The repository was not found. Either the name is wrong, or it " "is gated or private and therefore invisible without a token " "that can see it.") if EntryNotFoundError and isinstance(exc, EntryNotFoundError): return "The repository exists but the expected weight file is missing." text = str(exc) if "403" in text: return "Access was refused (HTTP 403), which usually means gating." if "401" in text: return "The request was unauthorized (HTTP 401), so no usable token was found." return f"The download failed with: {type(exc).__name__}: {text}" def _explain_backbone_failure(exc: Exception, *, name: str, repo: str, kwarg: str) -> RuntimeError: """Turn an opaque hub error into an actionable one.""" return RuntimeError(_ACCESS_HELP.format( name=name, repo=repo, reason=_diagnose(exc), kwarg=kwarg, leaf=repo.split("/")[-1])) class LumenConfig(PretrainedConfig): model_type = "lumen" def __init__(self, vision_repo: str = VISION_REPO, text_repo: str = TEXT_REPO, proj_dim: int = 512, vision_dim: int = 2560, text_dim: int = 768, lora_r: int = 4, lora_alpha: int = 8, lora_dropout: float = 0.1, logit_scale_init: float = 0.07, **kwargs): self.vision_repo = vision_repo self.text_repo = text_repo self.proj_dim = proj_dim # Virchow2 concatenates its class token with the mean of its patch # tokens, so the image feature is 2 x 1280. self.vision_dim = vision_dim self.text_dim = text_dim self.lora_r = lora_r self.lora_alpha = lora_alpha self.lora_dropout = lora_dropout self.logit_scale_init = logit_scale_init super().__init__(**kwargs) class _ContrastiveHead(nn.Module): def __init__(self, logit_scale_init: float = 0.07): super().__init__() self.logit_scale = nn.Parameter(torch.tensor(1 / logit_scale_init).log()) class LumenModel(PreTrainedModel): """Frozen Virchow2 and BioMedBERT joined by the released alignment.""" config_class = LumenConfig base_model_prefix = "lumen" main_input_name = "images" def __init__(self, config: LumenConfig, vision_path: str | None = None, text_path: str | None = None) -> None: super().__init__(config) from peft import LoraConfig, get_peft_model from transformers import AutoModel, AutoTokenizer try: import timm except ImportError as exc: # pragma: no cover - environment dependent raise RuntimeError( "Lumen needs `timm` to build the Virchow2 vision tower. " "Install it with: pip install timm" ) from exc vision_src = vision_path or f"hf-hub:{config.vision_repo}" try: vision = timm.create_model(vision_src, pretrained=True, mlp_layer=timm.layers.SwiGLUPacked, act_layer=torch.nn.SiLU) except Exception as exc: raise _explain_backbone_failure( exc, name="Virchow2 vision", repo=config.vision_repo, kwarg="vision_path") from exc text_src = text_path or config.text_repo try: text = AutoModel.from_pretrained(text_src) self._tokenizer = AutoTokenizer.from_pretrained(text_src) except Exception as exc: raise _explain_backbone_failure( exc, name="BioMedBERT text", repo=config.text_repo, kwarg="text_path") from exc common = dict(r=config.lora_r, lora_alpha=config.lora_alpha, lora_dropout=config.lora_dropout, bias="none") # Qualify with ``attn.``: PEFT matches a target as a dotted suffix, so a # bare "proj" also wraps ``patch_embed.proj``, which training never # adapted. That adds 64 untrained LoRA modules the released adapter has # no weights for, and the model still loads. self.img_encoder = get_peft_model( vision, LoraConfig(target_modules=["attn.qkv", "attn.proj"], **common)) self.txt_encoder = get_peft_model( text, LoraConfig(target_modules=["query", "key", "value", "attention.output.dense"], **common)) # A single linear map into the shared space, then layer norm. This is the # "direct" head of the training configuration. self.img_proj_net = nn.Sequential( nn.Linear(config.vision_dim, config.proj_dim), nn.LayerNorm(config.proj_dim)) self.txt_proj_net = nn.Sequential( nn.Linear(config.text_dim, config.proj_dim), nn.LayerNorm(config.proj_dim)) self.contrastive_head = _ContrastiveHead(config.logit_scale_init) @property def tokenizer(self): return self._tokenizer @property def logit_scale(self) -> torch.Tensor: return self.contrastive_head.logit_scale.exp().clamp(1.0, 20.0) def encode_image(self, images: torch.Tensor) -> torch.Tensor: out = self.img_encoder(images) # Class token concatenated with the mean of every remaining token. # # Virchow2 reserves tokens 1-4 as registers, and its model card pools # ``out[:, 5:]`` to exclude them. This model was trained pooling # ``out[:, 1:]``, so the registers are inside the mean. Reproducing the # released model means matching how it was trained, not how the backbone # is usually read: pooling from index 5 shifts the image embedding by # about 1.5e-3 per element, which no error surfaces and which cosine # similarity barely registers (0.99996). cls, rest = out[:, 0], out[:, 1:].mean(dim=1) return self.img_proj_net(torch.cat([cls, rest], dim=-1)) def encode_text(self, tokens: dict) -> torch.Tensor: out = self.txt_encoder(**tokens).last_hidden_state return self.txt_proj_net(out[:, 0]) def forward(self, images=None, tokens=None, **kwargs): # The arguments are ``images`` and ``tokens``, not the ``pixel_values`` # and ``input_ids`` most Hub models take. Silently ignoring the # conventional names returns a dict of Nones that only fails much later, # somewhere unrelated, so an unexpected keyword is refused here instead. if kwargs: raise TypeError( f"LumenModel.forward() got unexpected keyword(s) " f"{sorted(kwargs)}. It takes images= and tokens=." f" Pass either or both.") if images is None and tokens is None: raise TypeError( "LumenModel.forward() needs images=, tokens=, or both. " "With neither there is nothing to encode.") img = self.encode_image(images) if images is not None else None txt = self.encode_text(tokens) if tokens is not None else None return {"image_embeds": img, "text_embeds": txt, "logit_scale": self.logit_scale}