hf_compat.py

"""Optional Hugging Face compatibility with a fully functional local fallback.

The core model depends only on PyTorch.  When Transformers is installed the real
``PretrainedConfig``, ``PreTrainedModel``, ``GenerationMixin`` and ``ModelOutput``
classes are used.  The fallback exists so architecture tests and local research do
not silently depend on a network install.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, fields
from pathlib import Path
from typing import Any, ClassVar, Iterator, Mapping

import torch
from torch import nn

try:  # pragma: no cover - exercised only when Transformers is available.
    from transformers import GenerationConfig, PretrainedConfig, PreTrainedModel
    from transformers.generation import GenerationMixin
    from transformers.modeling_outputs import CausalLMOutputWithPast
    from transformers.utils import ModelOutput

    TRANSFORMERS_AVAILABLE = True

    class DendroGenerationConfig(GenerationConfig):
        """Keep Dendro's cache backend selector out of HF generation config.

        ``DendroOmniConfig.cache_implementation`` selects :class:`DendroKVCache`
        storage (including Dendro-specific values such as ``"int8"``).  Recent
        Transformers releases use the same field name for a different set of
        cache classes and reject those values while a model is being constructed.
        The model config retains the setting; only its generation-config copy is
        removed.
        """

        @classmethod
        def from_model_config(cls, model_config: Any) -> "DendroGenerationConfig":
            config_dict = model_config.to_dict() if not isinstance(model_config, dict) else dict(model_config)
            config_dict.pop("cache_implementation", None)
            return super().from_model_config(config_dict)
except Exception:  # pragma: no cover - fallback is covered instead.
    TRANSFORMERS_AVAILABLE = False

    class PretrainedConfig:
        """Small subset of the Hugging Face configuration contract."""

        model_type: ClassVar[str] = "model"

        def __init__(self, **kwargs: Any) -> None:
            for key, value in kwargs.items():
                setattr(self, key, value)

        @property
        def use_return_dict(self) -> bool:
            return bool(getattr(self, "return_dict", True))

        def to_dict(self) -> dict[str, Any]:
            result = dict(self.__dict__)
            result["model_type"] = self.model_type
            return result

        @classmethod
        def from_dict(cls, data: Mapping[str, Any], **kwargs: Any) -> "PretrainedConfig":
            merged = dict(data)
            merged.update(kwargs)
            merged.pop("model_type", None)
            return cls(**merged)

        def save_pretrained(self, save_directory: str | Path) -> None:
            path = Path(save_directory)
            path.mkdir(parents=True, exist_ok=True)
            (path / "config.json").write_text(
                json.dumps(self.to_dict(), indent=2, sort_keys=True), encoding="utf-8"
            )

        @classmethod
        def from_pretrained(cls, path: str | Path, **kwargs: Any) -> "PretrainedConfig":
            data = json.loads((Path(path) / "config.json").read_text(encoding="utf-8"))
            return cls.from_dict(data, **kwargs)

    class GenerationMixin:
        """Marker class used by the local model's own ``generate`` implementation."""

    class DendroGenerationConfig:
        """Fallback marker matching the Transformers generation config hook."""

    @dataclass
    class ModelOutput(Mapping[str, Any]):
        """Dataclass mapping behavior matching the useful part of HF ModelOutput."""

        def _items(self) -> list[tuple[str, Any]]:
            return [(field.name, getattr(self, field.name)) for field in fields(self) if getattr(self, field.name) is not None]

        def __getitem__(self, key: str | int | slice) -> Any:
            items = self._items()
            if isinstance(key, str):
                return dict(items)[key]
            return tuple(value for _name, value in items)[key]

        def __iter__(self) -> Iterator[str]:
            return (name for name, _value in self._items())

        def __len__(self) -> int:
            return len(self._items())

        def keys(self):  # type: ignore[override]
            return dict(self._items()).keys()

        def values(self):  # type: ignore[override]
            return dict(self._items()).values()

        def items(self):  # type: ignore[override]
            return dict(self._items()).items()

        def to_tuple(self) -> tuple[Any, ...]:
            return tuple(value for _name, value in self._items())

    @dataclass
    class CausalLMOutputWithPast(ModelOutput):
        loss: torch.Tensor | None = None
        logits: torch.Tensor | None = None
        past_key_values: Any = None
        hidden_states: tuple[torch.Tensor, ...] | None = None
        attentions: tuple[torch.Tensor, ...] | None = None

    class PreTrainedModel(nn.Module):
        """PyTorch-only persistence compatible with ``save_pretrained`` conventions."""

        config_class = PretrainedConfig
        base_model_prefix = "model"
        main_input_name = "input_ids"

        def __init__(self, config: PretrainedConfig, *args: Any, **kwargs: Any) -> None:
            del args, kwargs
            super().__init__()
            self.config = config

        def post_init(self) -> None:
            return None

        def save_pretrained(
            self,
            save_directory: str | Path,
            *,
            safe_serialization: bool = True,
            **_: Any,
        ) -> None:
            path = Path(save_directory)
            path.mkdir(parents=True, exist_ok=True)
            self.config.save_pretrained(path)
            state = self.state_dict()
            if safe_serialization:
                try:
                    from safetensors.torch import save_file

                    save_file(state, str(path / "model.safetensors"))
                    return
                except Exception:
                    pass
            torch.save(state, path / "pytorch_model.bin")

        @classmethod
        def from_pretrained(
            cls,
            pretrained_model_name_or_path: str | Path,
            *model_args: Any,
            config: PretrainedConfig | None = None,
            map_location: str | torch.device = "cpu",
            **kwargs: Any,
        ) -> "PreTrainedModel":
            path = Path(pretrained_model_name_or_path)
            if config is None:
                config = cls.config_class.from_pretrained(path)
            model = cls(config, *model_args, **kwargs)
            safe_path = path / "model.safetensors"
            torch_path = path / "pytorch_model.bin"
            if safe_path.exists():
                from safetensors.torch import load_file

                state = load_file(str(safe_path), device=str(map_location))
            elif torch_path.exists():
                state = torch.load(torch_path, map_location=map_location, weights_only=True)
            else:
                raise FileNotFoundError(f"No model.safetensors or pytorch_model.bin found in {path}")
            model.load_state_dict(state)
            return model


def is_transformers_available() -> bool:
    return TRANSFORMERS_AVAILABLE