File size: 3,198 Bytes
6f4e480 |
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 80 81 82 83 84 85 86 87 88 |
from __future__ import annotations
"""Remote-code modeling file for EXAONE-Path Slide/WSI encoder.
This file is imported by Transformers when using `trust_remote_code=True`.
Important:
- This file acts as a *thin AutoModel entrypoint*.
- The actual implementation lives in `exaonepath.models.slide_encoder_hf`.
- At runtime, the repository snapshot is downloaded via `snapshot_download`
and added to `sys.path` so that `exaonepath/` can be imported.
- Do NOT import sibling modules like `configuration_exaonepath_slide_encoder` here.
Transformers' remote-code dependency checker treats those imports as missing
third-party packages (e.g. it suggests `pip install configuration_exaonepath_slide_encoder`).
"""
from typing import Any, Dict, Optional
import importlib
import sys
from huggingface_hub import snapshot_download
from torch import Tensor, nn
from transformers import PretrainedConfig, PreTrainedModel
class ExaonePathSlideEncoderConfig(PretrainedConfig):
"""Self-contained Transformers config for EXAONE-Path Slide/WSI encoder.
Keep it here (in the modeling file) so we don't need a separate
`configuration_exaonepath_slide_encoder.py` on the Hub.
"""
model_type = "exaonepath_slide_encoder"
def __init__(self, wsi_cfg: Dict[str, Any] | None = None, **kwargs: Any):
self.wsi_cfg = dict(wsi_cfg or {})
super().__init__(**kwargs)
class ExaonePathSlideEncoderModel(PreTrainedModel):
config_class = ExaonePathSlideEncoderConfig
base_model_prefix = "slide_encoder"
def __init__(self, config: ExaonePathSlideEncoderConfig):
super().__init__(config)
# Ensure the repo code (including `exaonepath/`) is available at runtime.
# NOTE: config._name_or_path is usually the repo id when loaded from Hub.
repo_id = getattr(config, "_name_or_path", None) or getattr(config, "name_or_path", None)
if isinstance(repo_id, str) and repo_id:
local_root = snapshot_download(repo_id)
if local_root not in sys.path:
sys.path.insert(0, local_root)
WSIEncoder = getattr(importlib.import_module("exaonepath.models.slide_encoder_hf"), "WSIEncoder")
self.slide_encoder: nn.Module = WSIEncoder.from_wsi_config(wsi_cfg=config.wsi_cfg)
self.post_init()
def forward(
self,
patch_features: Tensor,
patch_mask: Tensor,
patch_coords: Optional[Tensor] = None,
patch_contour_index: Optional[Tensor] = None,
**kwargs: Any,
) -> Dict[str, Tensor]:
"""Return patch- and slide-level embeddings.
Returns a dict with exactly two keys:
- "patch_embedding": [B, N, C_in + D]
- "slide_embedding": [B, C_in + D]
Note: We intentionally return a plain dict (instead of a ModelOutput)
to make the remote-code API explicit and easy to use.
"""
out: Dict[str, Tensor] = self.slide_encoder(
patch_features=patch_features,
patch_mask=patch_mask,
patch_coords=patch_coords,
patch_contour_index=patch_contour_index,
)
return out
__all__ = ["ExaonePathSlideEncoderModel"]
|