Lumen / modeling_lumen.py

Rename the model to Lumen: registry key and checkpoint path in provenance

b89e805 verified about 2 months ago

11 kB

"""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=<float tensor "
                f"(B, 3, 224, 224)> and tokens=<output of model.tokenizer(...)>."
                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}

	"""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=<float tensor "
	f"(B, 3, 224, 224)> and tokens=<output of model.tokenizer(...)>."
	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}