luxical_hf_wrapper.py

from __future__ import annotations

from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence

import numpy as np
import pyarrow as pa
import torch
from torch import Tensor
from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import ModelOutput

from luxical.embedder import Embedder, _pack_int_dict, _unpack_int_dict
from luxical.sparse_to_dense_neural_nets import SparseToDenseEmbedder
from luxical.tokenization import ArrowTokenizer

DEFAULT_EMBEDDER_FILENAME = "luxical_one_embedder.npz"  # deprecated; no longer used


class LuxicalOneConfig(PretrainedConfig):
    """Configuration for the Luxical Huggingface wrapper.

    Generic for any Luxical `Embedder` serialized in format version 1.
    """

    model_type = "luxical-one"

    def __init__(
        self,
        *,
        max_ngram_length: int | None = None,
        embedding_dim: int | None = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.max_ngram_length = max_ngram_length
        self.embedding_dim = embedding_dim


@dataclass
class LuxicalOneModelOutput(ModelOutput):
    embeddings: Tensor


class LuxicalOneModel(PreTrainedModel):
    """Huggingface `PreTrainedModel` wrapper around a Luxical `Embedder`.

    Not tied to a specific checkpoint; reconstructs the `Embedder` from
    serialized state stored in the weights. Safetensors-only export.
    """
    config_class = LuxicalOneConfig

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):  # type: ignore[override]
        """Load model and reconstruct the Luxical embedder from safetensors.

        Keeps logic minimal and safetensors-only to avoid legacy branches.
        """
        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
        try:
            from transformers.utils import SAFE_WEIGHTS_NAME, cached_file
            from safetensors.torch import load_file as load_safetensors  # type: ignore
        except Exception:
            return model

        revision = kwargs.get("revision")
        cache_dir = kwargs.get("cache_dir")
        force_download = kwargs.get("force_download", False)
        proxies = kwargs.get("proxies")
        token = kwargs.get("token")
        local_files_only = kwargs.get("local_files_only", False)

        weight_path = None
        try:
            weight_path = cached_file(
                pretrained_model_name_or_path,
                SAFE_WEIGHTS_NAME,
                revision=revision,
                cache_dir=cache_dir,
                force_download=force_download,
                proxies=proxies,
                token=token,
                local_files_only=local_files_only,
            )
        except Exception:
            pass
        if weight_path is None:
            cand = Path(pretrained_model_name_or_path) / "model.safetensors"
            if cand.exists():
                weight_path = str(cand)

        if weight_path is not None:
            try:
                sd = load_safetensors(weight_path)
                model._embedder = _embedder_from_state_dict(sd)
                model._embedder_path = None
            except Exception:
                pass
        return model

    def __init__(
        self,
        config: LuxicalOneConfig,
        *,
        embedder: Embedder | None = None,
        embedder_path: str | Path | None = None,
    ) -> None:
        self._embedder: Embedder | None = embedder
        self._embedder_path: Path | None = (
            Path(embedder_path).resolve() if embedder_path is not None else None
        )
        super().__init__(config)

    def post_init(self) -> None:
        super().post_init()
        if self._embedder is not None:
            self.config.embedding_dim = self._embedder.embedding_dim
            self.config.max_ngram_length = self._embedder.max_ngram_length

    def forward(
        self,
        input_texts: Sequence[str] | pa.StringArray | None = None,
        *,
        batch_size: int = 4096,
        progress_bars: bool = False,
    ) -> LuxicalOneModelOutput:
        if input_texts is None:
            msg = "input_texts must be provided"
            raise ValueError(msg)
        embedder = self._ensure_embedder_loaded()
        embeddings_np = embedder(
            texts=input_texts,
            batch_size=batch_size,
            progress_bars=progress_bars,
        )
        embeddings = torch.from_numpy(embeddings_np)
        return LuxicalOneModelOutput(embeddings=embeddings)

    def save_pretrained(
        self,
        save_directory: str | Path,
        *args,
        **kwargs,
    ) -> tuple[OrderedDict[str, Tensor], LuxicalOneConfig]:
        save_path = Path(save_directory)
        save_path.mkdir(parents=True, exist_ok=True)
        # Prepare config with auto_map so AutoModel can import this module when
        # loading from a Hub/local repo with trust_remote_code=True.
        self.config.auto_map = {
            "AutoConfig": "luxical_hf_wrapper.LuxicalOneConfig",
            "AutoModel": "luxical_hf_wrapper.LuxicalOneModel",
        }
        # Persist the embedder inside a single Safetensors file.
        embedder = self._ensure_embedder_loaded()
        state_dict = _embedder_to_state_dict(embedder)
        from safetensors.torch import save_file as save_safetensors  # type: ignore
        save_safetensors(state_dict, str(save_path / "model.safetensors"))
        # Copy this module alongside to support remote code loading.
        import inspect
        import shutil

        module_src = Path(inspect.getsourcefile(LuxicalOneModel) or __file__).resolve()
        shutil.copyfile(module_src, save_path / "luxical_hf_wrapper.py")
        # Save config.json last.
        self.config.save_pretrained(save_path)
        return state_dict, self.config

    def load_state_dict(
        self, state_dict: OrderedDict[str, Tensor], strict: bool = True
    ):  # type: ignore[override]
        # Interpret the state dict as a serialized Luxical Embedder and rebuild it.
        try:
            self._embedder = _embedder_from_state_dict(state_dict)
            self._embedder_path = None
            # Update config fields if available
            self.config.embedding_dim = self._embedder.embedding_dim
            self.config.max_ngram_length = self._embedder.max_ngram_length
            return torch.nn.modules.module._IncompatibleKeys([], [])
        except KeyError:
            if strict:
                missing = list(state_dict.keys())
                raise NotImplementedError(
                    "LuxicalOneModel expected serialized embedder tensors; "
                    f"unexpected keys: {missing}"
                )
            return torch.nn.modules.module._IncompatibleKeys([], list(state_dict.keys()))

    def get_input_embeddings(self) -> torch.nn.Module:
        msg = "LuxicalOneModel does not expose token embeddings."
        raise NotImplementedError(msg)

    def set_input_embeddings(self, value: torch.nn.Module) -> None:
        msg = "LuxicalOneModel does not support replacing token embeddings."
        raise NotImplementedError(msg)

    def resize_token_embeddings(self, *args, **kwargs) -> None:
        msg = "LuxicalOneModel does not use token embeddings."
        raise NotImplementedError(msg)


    def _ensure_embedder_loaded(self) -> Embedder:
        if self._embedder is not None:
            return self._embedder
        raise RuntimeError(
            "Luxical embedder is not initialized. Load this model via "
            "AutoModel/LuxicalOneModel.from_pretrained so weights can be "
            "decoded into an Embedder."
        )

    # No legacy file-based loader; all state lives in model.safetensors.


def export_embedder_to_huggingface_directory(
    embedder: Embedder,
    save_directory: str | Path,
    *,
    config_overrides: dict[str, object] | None = None,
) -> Path:
    save_path = Path(save_directory)
    config = LuxicalOneConfig(
        max_ngram_length=embedder.max_ngram_length,
        embedding_dim=embedder.embedding_dim,
        **(config_overrides or {}),
    )
    config.name_or_path = str(save_path.resolve())
    model = LuxicalOneModel(config=config, embedder=embedder)
    model.save_pretrained(save_path)
    return save_path


# No global Auto* registration; exports include `auto_map` in config.json.


def _embedder_to_state_dict(embedder: Embedder) -> OrderedDict[str, Tensor]:
    sd: "OrderedDict[str, Tensor]" = OrderedDict()
    # Version
    sd["embedder.version"] = torch.tensor([1], dtype=torch.long)
    # Tokenizer json bytes
    tok_bytes = np.frombuffer(embedder.tokenizer.to_str().encode("utf-8"), dtype=np.uint8)
    sd["embedder.tokenizer"] = torch.from_numpy(tok_bytes.copy())
    # Recognized ngrams
    sd["embedder.recognized_ngrams"] = torch.from_numpy(embedder.recognized_ngrams.astype(np.int64, copy=False))
    # Hash map keys/values
    keys, vals = _unpack_int_dict(embedder.ngram_hash_to_ngram_idx)
    sd["embedder.ngram_keys"] = torch.from_numpy(keys.astype(np.int64, copy=False))
    sd["embedder.ngram_vals"] = torch.from_numpy(vals.astype(np.int64, copy=False))
    # IDF
    sd["embedder.idf_values"] = torch.from_numpy(embedder.idf_values.astype(np.float32, copy=False))
    # Layers
    layers = embedder.bow_to_dense_embedder.layers
    sd["embedder.num_layers"] = torch.tensor([len(layers)], dtype=torch.long)
    for i, layer in enumerate(layers):
        sd[f"embedder.nn_layer_{i}"] = torch.from_numpy(layer.astype(np.float32, copy=False))
    return sd


def _embedder_from_state_dict(state_dict: OrderedDict[str, Tensor]) -> Embedder:
    version = int(state_dict["embedder.version"][0].item())
    if version != 1:
        raise NotImplementedError(f"Unsupported embedder version: {version}")
    tok_bytes = bytes(state_dict["embedder.tokenizer"].cpu().numpy().astype(np.uint8).tolist())
    tokenizer = ArrowTokenizer(tok_bytes.decode("utf-8"))
    recognized_ngrams = state_dict["embedder.recognized_ngrams"].cpu().numpy().astype(np.int64, copy=False)
    keys = state_dict["embedder.ngram_keys"].cpu().numpy().astype(np.int64, copy=False)
    vals = state_dict["embedder.ngram_vals"].cpu().numpy().astype(np.int64, copy=False)
    ngram_map = _pack_int_dict(keys, vals)
    idf_values = state_dict["embedder.idf_values"].cpu().numpy().astype(np.float32, copy=False)
    num_layers = int(state_dict["embedder.num_layers"][0].item())
    layers = [
        state_dict[f"embedder.nn_layer_{i}"].cpu().numpy().astype(np.float32, copy=False)
        for i in range(num_layers)
    ]
    s2d = SparseToDenseEmbedder(layers=layers)
    return Embedder(
        tokenizer=tokenizer,
        recognized_ngrams=recognized_ngrams,
        ngram_hash_to_ngram_idx=ngram_map,
        idf_values=idf_values,
        bow_to_dense_embedder=s2d,
    )


def _parse_cli_args() -> tuple[str, dict[str, object]]:
    import argparse

    parser = argparse.ArgumentParser(
        description="Luxical One Huggingface wrapper: export and verify utilities.",
    )
    sub = parser.add_subparsers(dest="cmd", required=True)

    p_export = sub.add_parser(
        "export", help="Export a HF-formatted directory from a Luxical embedder .npz checkpoint"
    )
    p_export.add_argument(
        "--checkpoint",
        type=str,
        default=str(Path("/tmp/luxical_one_rc4.npz")),
        help="Path to Luxical embedder .npz checkpoint",
    )
    p_export.add_argument(
        "--output-dir",
        type=str,
        default=str(Path(__file__).resolve().parent / "artifacts" / "luxical_one_hf"),
        help="Directory to write the Huggingface-formatted model",
    )

    p_verify = sub.add_parser(
        "verify", help="Verify HF-loaded model matches native Embedder outputs"
    )
    p_verify.add_argument(
        "--checkpoint",
        type=str,
        default=str(Path("/tmp/luxical_one_rc4.npz")),
        help="Path to Luxical embedder .npz checkpoint",
    )
    p_verify.add_argument(
        "--export-dir",
        type=str,
        default=str(Path(__file__).resolve().parent / "artifacts" / "luxical_one_hf"),
        help="HF directory to create/use for verification",
    )
    p_verify.add_argument(
        "--batch-size", type=int, default=3, help="Batch size for verification"
    )

    args = parser.parse_args()
    return args.cmd, vars(args)


def _sample_texts() -> list[str]:
    return [
        "Luxical embeddings make tf-idf sparkle.",
        "This sentence tests the Huggingface wrapper path.",
        "Short.",
    ]


def _cmd_export(checkpoint: str, output_dir: str) -> None:
    ckpt_path = Path(checkpoint).expanduser().resolve()
    if not ckpt_path.exists():
        raise FileNotFoundError(
            f"Checkpoint not found at {ckpt_path}. Download with: aws s3 cp "
            "s3://datology-external-artifacts/luxical/luxical_one_rc4.npz "
            "/tmp/luxical_one_rc4.npz"
        )
    out_dir = Path(output_dir).expanduser().resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    embedder = Embedder.load(ckpt_path)
    export_embedder_to_huggingface_directory(embedder, out_dir)
    print(f"Huggingface directory written to {out_dir}")


def _cmd_verify(checkpoint: str, export_dir: str, batch_size: int) -> None:
    ckpt_path = Path(checkpoint).expanduser().resolve()
    if not ckpt_path.exists():
        raise FileNotFoundError(
            f"Checkpoint not found at {ckpt_path}. Download with: aws s3 cp "
            "s3://datology-external-artifacts/luxical/luxical_one_rc4.npz "
            "/tmp/luxical_one_rc4.npz"
        )
    exp_dir = Path(export_dir).expanduser().resolve()
    exp_dir.mkdir(parents=True, exist_ok=True)

    texts = _sample_texts()
    embedder = Embedder.load(ckpt_path)
    ref = embedder(texts, batch_size=batch_size)

    export_embedder_to_huggingface_directory(embedder, exp_dir)
    # Load using AutoModel so this mirrors user experience, with remote code.
    from transformers import AutoModel  # local import to keep top-level light
    model = AutoModel.from_pretrained(exp_dir, trust_remote_code=True)
    model.eval()
    with torch.inference_mode():
        out = (
            model(texts, batch_size=batch_size, progress_bars=False)
            .embeddings.cpu()
            .numpy()
        )
    import numpy as np

    np.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-6)
    print("Verification succeeded: Huggingface model matches embedder output.")


if __name__ == "__main__":
    cmd, kv = _parse_cli_args()
    if cmd == "export":
        _cmd_export(checkpoint=str(kv["checkpoint"]), output_dir=str(kv["output_dir"]))
    elif cmd == "verify":
        _cmd_verify(
            checkpoint=str(kv["checkpoint"]),
            export_dir=str(kv["export_dir"]),
            batch_size=int(kv["batch_size"]),
        )
    else:
        raise SystemExit(f"Unknown command: {cmd}")