jolia_shim.py

"""Minimal stand-ins for the ``raidium.rd.hub`` base classes.

The vendored Atlas modules (``jolia_multimodal_atlas`` etc.) were copied
verbatim from the internal ``raidium.rd.models`` library, which subclasses
``raidium.rd.hub.utils.{config.Config, base_model.BaseModel,
base_preprocessing.BasePreprocessing}``. Those base classes carry a lot of
hub-registry machinery the public Hugging Face release does not need, so we
replace them with the three tiny shims below. Behaviour relevant to inference
(config field defaults, ``self.config`` storage, ``nn.Module`` semantics) is
preserved exactly.
"""

from __future__ import annotations

import torch.nn as nn


class Config:
    """Lightweight replacement for ``raidium.rd.hub.utils.config.Config``.

    Sets every field annotated anywhere in the MRO from the matching keyword
    argument, falling back to the class-level default. This reproduces the
    pydantic-style ``MyConfig(field=value)`` construction the vendored Atlas
    configs rely on, without the pydantic dependency.
    """

    def __init__(self, **kwargs: object) -> None:
        fields: dict[str, object] = {}
        for klass in reversed(type(self).__mro__):
            fields.update(getattr(klass, "__annotations__", {}))
        for name in fields:
            setattr(self, name, kwargs.get(name, getattr(type(self), name, None)))


class BaseModel(nn.Module):
    """Replacement for ``raidium.rd.hub.utils.base_model.BaseModel``.

    Only the constructor contract used by the vendored backbone is kept: store
    ``config`` and behave as a plain ``nn.Module``. Hub push/pull methods are
    intentionally absent — Hugging Face's ``PreTrainedModel`` owns that role.
    """

    config_class = None

    def __init__(self, config: Config) -> None:
        super().__init__()
        self.config = config


class BasePreprocessing(nn.Module):
    """Replacement for ``raidium.rd.hub.utils.base_preprocessing.BasePreprocessing``."""

    config_class = None
    transform_input = "metadata"

    def __init__(self, config: Config) -> None:
        super().__init__()
        self.config = config