Skip to content

Teachers

Teachers can be trained inside RecDistillery or imported from external artifacts. In both cases, the runtime representation is a framework-neutral TeacherState saved as a .teacher artifact.

Teacher Training

Train a native teacher with:

python scripts/teacher_training/teacher_training.py \
  --framework recbole \
  --model BPRMF \
  --dataset citeulike

Complete experiment configs are stored in:

config/experiments/teacher/

Teacher Import

External teachers are converted through scripts/recdistill/import_teacher.py. The import system supports generic checkpoints, prediction JSON exports, and RecBole .pth checkpoints.

python scripts/recdistill/import_teacher.py --list-adapters

Teacher State

TeacherState is the normalized runtime object used after training or import. It hides framework-specific checkpoint formats and gives distillers a single interface.

Object Represents Key attributes
TeacherState A trained or imported teacher. user_embeddings, item_embeddings, metadata, scorer
PrecomputedScoresScorer A dense user-item score matrix. scores
PrecomputedTopKScorer A sparse top-k ranking export. topk_items, topk_scores, fill_value, num_items_override
TeacherScorer Protocol for scorer-only teachers. to, score_items_for_user

TeacherScorer

Bases: Protocol

Protocol implemented by scorer-only teacher representations.

A scorer computes item scores for one user without requiring explicit user/item embedding matrices. Import adapters use it for prediction exports, top-k rankings, or dense score matrices.

Source code in recdistill/teachers/state.py
class TeacherScorer(Protocol):
    """Protocol implemented by scorer-only teacher representations.

    A scorer computes item scores for one user without requiring explicit
    user/item embedding matrices. Import adapters use it for prediction exports,
    top-k rankings, or dense score matrices.
    """

    def to(self, device: torch.device | str): ...
    def score_items_for_user(self, user: int, num_items: int) -> torch.Tensor: ...

to(device: torch.device | str)

Source code in recdistill/teachers/state.py
def to(self, device: torch.device | str): ...

score_items_for_user(user: int, num_items: int) -> torch.Tensor

Source code in recdistill/teachers/state.py
def score_items_for_user(self, user: int, num_items: int) -> torch.Tensor: ...

PrecomputedScoresScorer dataclass

Dense precomputed teacher score matrix.

Attributes:

Name Type Description
scores Tensor

Tensor with shape [num_users, num_items]. Each row contains the teacher scores for all candidate items of one user.

Source code in recdistill/teachers/state.py
@dataclass
class PrecomputedScoresScorer:
    """Dense precomputed teacher score matrix.

    Attributes:
        scores: Tensor with shape `[num_users, num_items]`. Each row contains
            the teacher scores for all candidate items of one user.
    """

    scores: torch.Tensor

    def __post_init__(self) -> None:
        self.scores = torch.as_tensor(self.scores, dtype=torch.float32)
        if self.scores.ndim != 2:
            raise ValueError("Precomputed score matrix must be a rank-2 tensor.")

    @property
    def num_users(self) -> int:
        return int(self.scores.size(0))

    @property
    def num_items(self) -> int:
        return int(self.scores.size(1))

    def to(self, device: torch.device | str) -> "PrecomputedScoresScorer":
        return PrecomputedScoresScorer(scores=self.scores.to(device))

    def score_items_for_user(self, user: int, num_items: int) -> torch.Tensor:
        """Return the score vector for a user, padded or truncated to `num_items`."""
        if user < 0 or user >= self.num_users:
            raise IndexError(f"User index out of bounds for precomputed scores: {user}")
        scores = self.scores[user]
        if scores.numel() >= num_items:
            return scores[:num_items]
        padded = scores.new_full((num_items,), float("-inf"))
        padded[: scores.numel()] = scores
        return padded

scores: torch.Tensor instance-attribute

num_users: int property

num_items: int property

__init__(scores: torch.Tensor) -> None

__post_init__() -> None

Source code in recdistill/teachers/state.py
def __post_init__(self) -> None:
    self.scores = torch.as_tensor(self.scores, dtype=torch.float32)
    if self.scores.ndim != 2:
        raise ValueError("Precomputed score matrix must be a rank-2 tensor.")

to(device: torch.device | str) -> 'PrecomputedScoresScorer'

Source code in recdistill/teachers/state.py
def to(self, device: torch.device | str) -> "PrecomputedScoresScorer":
    return PrecomputedScoresScorer(scores=self.scores.to(device))

score_items_for_user(user: int, num_items: int) -> torch.Tensor

Return the score vector for a user, padded or truncated to num_items.

Source code in recdistill/teachers/state.py
def score_items_for_user(self, user: int, num_items: int) -> torch.Tensor:
    """Return the score vector for a user, padded or truncated to `num_items`."""
    if user < 0 or user >= self.num_users:
        raise IndexError(f"User index out of bounds for precomputed scores: {user}")
    scores = self.scores[user]
    if scores.numel() >= num_items:
        return scores[:num_items]
    padded = scores.new_full((num_items,), float("-inf"))
    padded[: scores.numel()] = scores
    return padded

PrecomputedTopKScorer dataclass

Sparse scorer backed by precomputed ranked items.

Attributes:

Name Type Description
topk_items Tensor

Integer tensor with shape [num_users, top_k].

topk_scores Tensor | None

Optional score tensor aligned with topk_items.

fill_value float

Score assigned to items that are absent from the top-k list.

num_items_override int | None

Optional catalog size when it cannot be inferred from the maximum item id.

Source code in recdistill/teachers/state.py
@dataclass
class PrecomputedTopKScorer:
    """Sparse scorer backed by precomputed ranked items.

    Attributes:
        topk_items: Integer tensor with shape `[num_users, top_k]`.
        topk_scores: Optional score tensor aligned with `topk_items`.
        fill_value: Score assigned to items that are absent from the top-k list.
        num_items_override: Optional catalog size when it cannot be inferred
            from the maximum item id.
    """

    topk_items: torch.Tensor
    topk_scores: torch.Tensor | None = None
    fill_value: float = float("-inf")
    num_items_override: int | None = None

    def __post_init__(self) -> None:
        self.topk_items = torch.as_tensor(self.topk_items, dtype=torch.long)
        if self.topk_items.ndim != 2:
            raise ValueError("Precomputed top-k items must be a rank-2 tensor.")
        if self.topk_scores is not None:
            self.topk_scores = torch.as_tensor(self.topk_scores, dtype=torch.float32)
            if self.topk_scores.shape != self.topk_items.shape:
                raise ValueError("Precomputed top-k scores must have the same shape as top-k items.")

    @property
    def num_users(self) -> int:
        return int(self.topk_items.size(0))

    @property
    def top_k(self) -> int:
        return int(self.topk_items.size(1))

    @property
    def num_items(self) -> int:
        if self.num_items_override is not None:
            return int(self.num_items_override)
        valid = self.topk_items[self.topk_items >= 0]
        return int(valid.max().item() + 1) if valid.numel() else 0

    def to(self, device: torch.device | str) -> "PrecomputedTopKScorer":
        return PrecomputedTopKScorer(
            topk_items=self.topk_items.to(device),
            topk_scores=self.topk_scores.to(device) if self.topk_scores is not None else None,
            fill_value=float(self.fill_value),
            num_items_override=self.num_items_override,
        )

    def score_items_for_user(self, user: int, num_items: int) -> torch.Tensor:
        """Expand one user's top-k ranking into a full score vector."""
        if user < 0 or user >= self.num_users:
            raise IndexError(f"User index out of bounds for precomputed top-k: {user}")
        scores = torch.full((num_items,), float(self.fill_value), dtype=torch.float32, device=self.topk_items.device)
        items = self.topk_items[user]
        valid = (items >= 0) & (items < num_items)
        if not bool(valid.any()):
            return scores
        valid_items = items[valid]
        if self.topk_scores is not None:
            valid_scores = self.topk_scores[user][valid].to(dtype=torch.float32, device=scores.device)
        else:
            ranks = torch.arange(valid_items.numel(), dtype=torch.float32, device=scores.device)
            valid_scores = -ranks
        scores[valid_items] = valid_scores
        return scores

topk_items: torch.Tensor instance-attribute

topk_scores: torch.Tensor | None = None class-attribute instance-attribute

fill_value: float = float('-inf') class-attribute instance-attribute

num_items_override: int | None = None class-attribute instance-attribute

num_users: int property

top_k: int property

num_items: int property

__init__(topk_items: torch.Tensor, topk_scores: torch.Tensor | None = None, fill_value: float = float('-inf'), num_items_override: int | None = None) -> None

__post_init__() -> None

Source code in recdistill/teachers/state.py
def __post_init__(self) -> None:
    self.topk_items = torch.as_tensor(self.topk_items, dtype=torch.long)
    if self.topk_items.ndim != 2:
        raise ValueError("Precomputed top-k items must be a rank-2 tensor.")
    if self.topk_scores is not None:
        self.topk_scores = torch.as_tensor(self.topk_scores, dtype=torch.float32)
        if self.topk_scores.shape != self.topk_items.shape:
            raise ValueError("Precomputed top-k scores must have the same shape as top-k items.")

to(device: torch.device | str) -> 'PrecomputedTopKScorer'

Source code in recdistill/teachers/state.py
def to(self, device: torch.device | str) -> "PrecomputedTopKScorer":
    return PrecomputedTopKScorer(
        topk_items=self.topk_items.to(device),
        topk_scores=self.topk_scores.to(device) if self.topk_scores is not None else None,
        fill_value=float(self.fill_value),
        num_items_override=self.num_items_override,
    )

score_items_for_user(user: int, num_items: int) -> torch.Tensor

Expand one user's top-k ranking into a full score vector.

Source code in recdistill/teachers/state.py
def score_items_for_user(self, user: int, num_items: int) -> torch.Tensor:
    """Expand one user's top-k ranking into a full score vector."""
    if user < 0 or user >= self.num_users:
        raise IndexError(f"User index out of bounds for precomputed top-k: {user}")
    scores = torch.full((num_items,), float(self.fill_value), dtype=torch.float32, device=self.topk_items.device)
    items = self.topk_items[user]
    valid = (items >= 0) & (items < num_items)
    if not bool(valid.any()):
        return scores
    valid_items = items[valid]
    if self.topk_scores is not None:
        valid_scores = self.topk_scores[user][valid].to(dtype=torch.float32, device=scores.device)
    else:
        ranks = torch.arange(valid_items.numel(), dtype=torch.float32, device=scores.device)
        valid_scores = -ranks
    scores[valid_items] = valid_scores
    return scores

TeacherState dataclass

Framework-neutral teacher representation used by distillation.

A teacher can be represented either by user/item embeddings or by a scorer. The same state object is used for native teachers, imported checkpoints, prediction JSON files, and serialized .teacher artifacts.

Attributes:

Name Type Description
user_embeddings Tensor | None

Optional user embedding matrix.

item_embeddings Tensor | None

Optional item embedding matrix.

metadata dict[str, object]

Free-form provenance and mapping information.

scorer TeacherScorer | None

Optional scorer-only representation.

Source code in recdistill/teachers/state.py
@dataclass
class TeacherState:
    """Framework-neutral teacher representation used by distillation.

    A teacher can be represented either by user/item embeddings or by a scorer.
    The same state object is used for native teachers, imported checkpoints,
    prediction JSON files, and serialized `.teacher` artifacts.

    Attributes:
        user_embeddings: Optional user embedding matrix.
        item_embeddings: Optional item embedding matrix.
        metadata: Free-form provenance and mapping information.
        scorer: Optional scorer-only representation.
    """

    user_embeddings: torch.Tensor | None = None
    item_embeddings: torch.Tensor | None = None
    metadata: dict[str, object] = field(default_factory=dict)
    scorer: TeacherScorer | None = None

    def __post_init__(self) -> None:
        if self.user_embeddings is not None:
            self.user_embeddings = torch.as_tensor(self.user_embeddings, dtype=torch.float32)
            if self.user_embeddings.ndim != 2:
                raise ValueError("TeacherState.user_embeddings must be a rank-2 tensor.")
        if self.item_embeddings is not None:
            self.item_embeddings = torch.as_tensor(self.item_embeddings, dtype=torch.float32)
            if self.item_embeddings.ndim != 2:
                raise ValueError("TeacherState.item_embeddings must be a rank-2 tensor.")
        if (self.user_embeddings is None) != (self.item_embeddings is None):
            raise ValueError("TeacherState requires both user and item embeddings, or neither.")
        if self.user_embeddings is not None and self.item_embeddings is not None:
            if self.user_embeddings.size(1) != self.item_embeddings.size(1):
                raise ValueError("TeacherState user/item embeddings must share the same embedding dimension.")
        if self.user_embeddings is None and self.scorer is None:
            raise ValueError("TeacherState requires embeddings or a scorer.")

    @property
    def num_users(self) -> int:
        if self.user_embeddings is not None:
            return int(self.user_embeddings.size(0))
        if hasattr(self.scorer, "num_users"):
            return int(getattr(self.scorer, "num_users"))
        if "num_users" in self.metadata:
            return int(self.metadata["num_users"])
        raise ValueError("TeacherState.num_users is unavailable without embeddings, scorer metadata, or num_users metadata.")

    @property
    def num_items(self) -> int:
        if self.item_embeddings is not None:
            return int(self.item_embeddings.size(0))
        if hasattr(self.scorer, "num_items"):
            return int(getattr(self.scorer, "num_items"))
        if "num_items" in self.metadata:
            return int(self.metadata["num_items"])
        raise ValueError("TeacherState.num_items is unavailable without embeddings, scorer metadata, or num_items metadata.")

    @property
    def embedding_dim(self) -> int:
        if self.user_embeddings is None:
            raise ValueError("TeacherState has no embedding representation.")
        return int(self.user_embeddings.size(1))

    @property
    def device(self) -> torch.device:
        if self.user_embeddings is not None:
            return self.user_embeddings.device
        if isinstance(self.scorer, PrecomputedScoresScorer):
            return self.scorer.scores.device
        if isinstance(self.scorer, PrecomputedTopKScorer):
            return self.scorer.topk_items.device
        return torch.device("cpu")

    @property
    def has_embeddings(self) -> bool:
        return self.user_embeddings is not None and self.item_embeddings is not None

    def to(self, device: torch.device | str) -> "TeacherState":
        """Return a copy of the teacher state moved to `device`."""
        scorer = self.scorer.to(device) if self.scorer is not None else None
        return TeacherState(
            user_embeddings=self.user_embeddings.to(device) if self.user_embeddings is not None else None,
            item_embeddings=self.item_embeddings.to(device) if self.item_embeddings is not None else None,
            metadata=dict(self.metadata),
            scorer=scorer,
        )

user_embeddings: torch.Tensor | None = None class-attribute instance-attribute

item_embeddings: torch.Tensor | None = None class-attribute instance-attribute

metadata: dict[str, object] = field(default_factory=dict) class-attribute instance-attribute

scorer: TeacherScorer | None = None class-attribute instance-attribute

num_users: int property

num_items: int property

embedding_dim: int property

device: torch.device property

has_embeddings: bool property

__init__(user_embeddings: torch.Tensor | None = None, item_embeddings: torch.Tensor | None = None, metadata: dict[str, object] = dict(), scorer: TeacherScorer | None = None) -> None

__post_init__() -> None

Source code in recdistill/teachers/state.py
def __post_init__(self) -> None:
    if self.user_embeddings is not None:
        self.user_embeddings = torch.as_tensor(self.user_embeddings, dtype=torch.float32)
        if self.user_embeddings.ndim != 2:
            raise ValueError("TeacherState.user_embeddings must be a rank-2 tensor.")
    if self.item_embeddings is not None:
        self.item_embeddings = torch.as_tensor(self.item_embeddings, dtype=torch.float32)
        if self.item_embeddings.ndim != 2:
            raise ValueError("TeacherState.item_embeddings must be a rank-2 tensor.")
    if (self.user_embeddings is None) != (self.item_embeddings is None):
        raise ValueError("TeacherState requires both user and item embeddings, or neither.")
    if self.user_embeddings is not None and self.item_embeddings is not None:
        if self.user_embeddings.size(1) != self.item_embeddings.size(1):
            raise ValueError("TeacherState user/item embeddings must share the same embedding dimension.")
    if self.user_embeddings is None and self.scorer is None:
        raise ValueError("TeacherState requires embeddings or a scorer.")

to(device: torch.device | str) -> 'TeacherState'

Return a copy of the teacher state moved to device.

Source code in recdistill/teachers/state.py
def to(self, device: torch.device | str) -> "TeacherState":
    """Return a copy of the teacher state moved to `device`."""
    scorer = self.scorer.to(device) if self.scorer is not None else None
    return TeacherState(
        user_embeddings=self.user_embeddings.to(device) if self.user_embeddings is not None else None,
        item_embeddings=self.item_embeddings.to(device) if self.item_embeddings is not None else None,
        metadata=dict(self.metadata),
        scorer=scorer,
    )

Teacher Sources

TeacherSource describes where an external teacher comes from and which hints the adapter registry can use during import.

Attribute Meaning
path Main artifact path, such as .teacher, .pth, .pt, .ckpt, or .json.
framework Framework hint used for adapter resolution.
format File/representation hint used for adapter resolution.
model_name Optional model name saved into teacher metadata.
adapter Explicit custom adapter import path.
metadata Dataset, id mapping, provenance, and any additional import metadata.

TeacherSource dataclass

Input descriptor consumed by teacher import adapters.

Attributes:

Name Type Description
path Path | None

Main checkpoint, prediction, or .teacher artifact path.

framework str

Framework hint such as recbole, elliot, or external.

format str

Format hint such as checkpoint, predictions_json, or recbole_pth.

model_name str | None

Optional model/backbone name stored in metadata.

adapter str | None

Optional explicit adapter import path.

user_embeddings_path Path | None

Optional external user embedding file.

item_embeddings_path Path | None

Optional external item embedding file.

score_matrix_path Path | None

Optional dense score matrix file.

topk_items_path Path | None

Optional top-k item matrix file.

topk_scores_path Path | None

Optional top-k score matrix file.

metadata dict[str, Any]

Extra provenance, mapping, and dataset information.

Source code in recdistill/teachers/source.py
@dataclass(frozen=True)
class TeacherSource:
    """Input descriptor consumed by teacher import adapters.

    Attributes:
        path: Main checkpoint, prediction, or `.teacher` artifact path.
        framework: Framework hint such as `recbole`, `elliot`, or `external`.
        format: Format hint such as `checkpoint`, `predictions_json`, or
            `recbole_pth`.
        model_name: Optional model/backbone name stored in metadata.
        adapter: Optional explicit adapter import path.
        user_embeddings_path: Optional external user embedding file.
        item_embeddings_path: Optional external item embedding file.
        score_matrix_path: Optional dense score matrix file.
        topk_items_path: Optional top-k item matrix file.
        topk_scores_path: Optional top-k score matrix file.
        metadata: Extra provenance, mapping, and dataset information.
    """

    path: Path | None = None
    framework: str = "auto"
    format: str = "auto"
    model_name: str | None = None
    adapter: str | None = None
    user_embeddings_path: Path | None = None
    item_embeddings_path: Path | None = None
    score_matrix_path: Path | None = None
    topk_items_path: Path | None = None
    topk_scores_path: Path | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_path(
        cls,
        path: str | Path,
        *,
        framework: str = "auto",
        format: str = "auto",
        model_name: str | None = None,
        adapter: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> "TeacherSource":
        """Create a source descriptor from one primary artifact path."""
        return cls(
            path=Path(path),
            framework=framework,
            format=format,
            model_name=model_name,
            adapter=adapter,
            metadata=metadata or {},
        )

path: Path | None = None class-attribute instance-attribute

framework: str = 'auto' class-attribute instance-attribute

format: str = 'auto' class-attribute instance-attribute

model_name: str | None = None class-attribute instance-attribute

adapter: str | None = None class-attribute instance-attribute

user_embeddings_path: Path | None = None class-attribute instance-attribute

item_embeddings_path: Path | None = None class-attribute instance-attribute

score_matrix_path: Path | None = None class-attribute instance-attribute

topk_items_path: Path | None = None class-attribute instance-attribute

topk_scores_path: Path | None = None class-attribute instance-attribute

metadata: dict[str, Any] = field(default_factory=dict) class-attribute instance-attribute

__init__(path: Path | None = None, framework: str = 'auto', format: str = 'auto', model_name: str | None = None, adapter: str | None = None, user_embeddings_path: Path | None = None, item_embeddings_path: Path | None = None, score_matrix_path: Path | None = None, topk_items_path: Path | None = None, topk_scores_path: Path | None = None, metadata: dict[str, Any] = dict()) -> None

from_path(path: str | Path, *, framework: str = 'auto', format: str = 'auto', model_name: str | None = None, adapter: str | None = None, metadata: dict[str, Any] | None = None) -> 'TeacherSource' classmethod

Create a source descriptor from one primary artifact path.

Source code in recdistill/teachers/source.py
@classmethod
def from_path(
    cls,
    path: str | Path,
    *,
    framework: str = "auto",
    format: str = "auto",
    model_name: str | None = None,
    adapter: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> "TeacherSource":
    """Create a source descriptor from one primary artifact path."""
    return cls(
        path=Path(path),
        framework=framework,
        format=format,
        model_name=model_name,
        adapter=adapter,
        metadata=metadata or {},
    )

Loading And Serialization

load_teacher(source: TeacherSource | str | Path, device: torch.device | str | None = None) -> TeacherState

Source code in recdistill/teachers/loaders.py
def load_teacher(
    source: TeacherSource | str | Path,
    device: torch.device | str | None = None,
) -> TeacherState:
    if not isinstance(source, TeacherSource):
        source = TeacherSource.from_path(source)
    return _load_teacher_state_from_source(source, device=device)

load_teacher_state(source: TeacherSource | str | Path, device: torch.device | str | None = None) -> TeacherState

Load a teacher state from any registered RecDistill teacher source.

This is the preferred public loader. It accepts either a TeacherSource or a direct artifact path and dispatches through the teacher adapter registry.

Source code in recdistill/teachers/loaders.py
def load_teacher_state(
    source: TeacherSource | str | Path,
    device: torch.device | str | None = None,
) -> TeacherState:
    """Load a teacher state from any registered RecDistill teacher source.

    This is the preferred public loader. It accepts either a TeacherSource or a
    direct artifact path and dispatches through the teacher adapter registry.
    """
    return load_teacher(source, device=device)

register_default_teacher_adapters() -> None

Source code in recdistill/teachers/loaders.py
def register_default_teacher_adapters() -> None:
    from recdistill.teachers.adapters import (
        CheckpointAdapter,
        PredictionsJsonAdapter,
        RecBolePthAdapter,
    )

    register_teacher_adapter(CheckpointAdapter(), "torch", "torch_checkpoint", "teacher")
    register_teacher_adapter(PredictionsJsonAdapter(), "prediction_json", "json_predictions", "json")
    register_teacher_adapter(RecBolePthAdapter(), "recbole", "pth")

TEACHER_FORMAT_VERSION = 'recdistill.teacher.v2' module-attribute

SUPPORTED_TEACHER_FORMAT_VERSIONS = {'recdistill.teacher.v1', TEACHER_FORMAT_VERSION} module-attribute

teacher_state_to_payload(state: TeacherState, *, framework: str | None = None, model_name: str | None = None, metadata: dict[str, Any] | None = None) -> dict[str, Any]

Source code in recdistill/teachers/serialization.py
def teacher_state_to_payload(
    state: TeacherState,
    *,
    framework: str | None = None,
    model_name: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> dict[str, Any]:
    merged_metadata = dict(state.metadata)
    if metadata:
        merged_metadata.update(metadata)
    if framework is not None:
        merged_metadata["framework"] = framework
    if model_name is not None:
        merged_metadata["model_name"] = model_name

    scorer_payload = _scorer_to_payload(state)
    if scorer_payload is None and not state.has_embeddings:
        raise ValueError("Cannot serialize a non-embedding teacher with an unsupported scorer type.")

    return {
        "format_version": TEACHER_FORMAT_VERSION,
        "created_at_utc": utc_now_iso(),
        "framework": framework or merged_metadata.get("framework") or merged_metadata.get("source"),
        "model_name": model_name or merged_metadata.get("model_name"),
        "num_users": state.num_users,
        "num_items": state.num_items,
        "user_embeddings": state.user_embeddings.detach().cpu() if state.user_embeddings is not None else None,
        "item_embeddings": state.item_embeddings.detach().cpu() if state.item_embeddings is not None else None,
        "scorer": scorer_payload,
        "local_to_public_user_id": merged_metadata.get("local_to_public_user_id"),
        "local_to_public_item_id": merged_metadata.get("local_to_public_item_id"),
        "public_to_local_user_id": merged_metadata.get("public_to_local_user_id"),
        "public_to_local_item_id": merged_metadata.get("public_to_local_item_id"),
        "metadata": merged_metadata,
    }

teacher_state_from_payload(payload: dict[str, Any]) -> TeacherState

Source code in recdistill/teachers/serialization.py
def teacher_state_from_payload(payload: dict[str, Any]) -> TeacherState:
    if payload.get("format_version") not in SUPPORTED_TEACHER_FORMAT_VERSIONS:
        raise ValueError(
            f"Unsupported teacher format: {payload.get('format_version')!r}. "
            f"Expected one of {sorted(SUPPORTED_TEACHER_FORMAT_VERSIONS)}."
        )
    metadata = dict(payload.get("metadata") or {})
    for key in (
        "framework",
        "model_name",
        "created_at_utc",
        "num_users",
        "num_items",
        "local_to_public_user_id",
        "local_to_public_item_id",
        "public_to_local_user_id",
        "public_to_local_item_id",
    ):
        if key in payload and payload[key] is not None:
            metadata.setdefault(key, payload[key])
    metadata.setdefault("source", "checkpoint_teacher")
    return TeacherState(
        user_embeddings=(
            torch.as_tensor(payload["user_embeddings"], dtype=torch.float32)
            if payload.get("user_embeddings") is not None
            else None
        ),
        item_embeddings=(
            torch.as_tensor(payload["item_embeddings"], dtype=torch.float32)
            if payload.get("item_embeddings") is not None
            else None
        ),
        metadata=metadata,
        scorer=_scorer_from_payload(payload),
    )

save_teacher_state(path: str | Path, state: TeacherState, *, framework: str | None = None, model_name: str | None = None, metadata: dict[str, Any] | None = None) -> dict[str, Any]

Source code in recdistill/teachers/serialization.py
def save_teacher_state(
    path: str | Path,
    state: TeacherState,
    *,
    framework: str | None = None,
    model_name: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> dict[str, Any]:
    output_path = Path(path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    payload = teacher_state_to_payload(
        state,
        framework=framework,
        model_name=model_name,
        metadata=metadata,
    )
    torch.save(payload, output_path)
    return payload

load_teacher_payload(path: str | Path) -> dict[str, Any]

Source code in recdistill/teachers/serialization.py
def load_teacher_payload(path: str | Path) -> dict[str, Any]:
    payload = torch.load(Path(path), map_location="cpu", weights_only=False)
    if not isinstance(payload, dict):
        raise TypeError(f"Unsupported teacher payload type: {type(payload)!r}")
    return payload

Adapter Registry

The registry maps teacher sources to import adapters. Callers can register new adapters, list available keys, resolve the adapter for a source, or load a TeacherState directly.

Function Purpose
register_teacher_adapter Adds an adapter and optional aliases.
available_teacher_adapters Lists registered adapter keys.
resolve_teacher_adapter Selects the adapter that can load a TeacherSource.
load_teacher_state Resolves and loads a TeacherState.

TeacherAdapter

Bases: Protocol

Protocol implemented by teacher import adapters.

Attributes:

Name Type Description
name str

Stable registry key used by --list-adapters and adapter resolution.

Source code in recdistill/teachers/registry.py
class TeacherAdapter(Protocol):
    """Protocol implemented by teacher import adapters.

    Attributes:
        name: Stable registry key used by `--list-adapters` and adapter
            resolution.
    """

    name: str

    def can_load(self, source: TeacherSource) -> bool: ...

    def load(self, source: TeacherSource, device: torch.device | str | None = None) -> TeacherState: ...

name: str instance-attribute

can_load(source: TeacherSource) -> bool

Source code in recdistill/teachers/registry.py
def can_load(self, source: TeacherSource) -> bool: ...

load(source: TeacherSource, device: torch.device | str | None = None) -> TeacherState

Source code in recdistill/teachers/registry.py
def load(self, source: TeacherSource, device: torch.device | str | None = None) -> TeacherState: ...

register_teacher_adapter(adapter: TeacherAdapter, *aliases: str) -> None

Register an adapter under its primary name and optional aliases.

Source code in recdistill/teachers/registry.py
def register_teacher_adapter(adapter: TeacherAdapter, *aliases: str) -> None:
    """Register an adapter under its primary name and optional aliases."""
    names = (adapter.name, *aliases)
    for name in names:
        _ADAPTERS[_normalize(name)] = adapter

available_teacher_adapters() -> tuple[str, ...]

Return the sorted adapter keys currently available in the registry.

Source code in recdistill/teachers/registry.py
def available_teacher_adapters() -> tuple[str, ...]:
    """Return the sorted adapter keys currently available in the registry."""
    return tuple(sorted(_ADAPTERS))

resolve_teacher_adapter(source: TeacherSource) -> TeacherAdapter

Choose the adapter able to load source.

Resolution first honors an explicit adapter import path, then exact format or framework hints, and finally asks registered adapters whether they can load the source.

Source code in recdistill/teachers/registry.py
def resolve_teacher_adapter(source: TeacherSource) -> TeacherAdapter:
    """Choose the adapter able to load `source`.

    Resolution first honors an explicit adapter import path, then exact format
    or framework hints, and finally asks registered adapters whether they can
    load the source.
    """
    if source.adapter:
        return _load_adapter_object(source.adapter)

    checkpoint_adapter = _ADAPTERS.get("checkpoint") or _ADAPTERS.get("teacher")
    if checkpoint_adapter is not None and source.path is not None:
        try:
            suffix = source.path.suffix.lower()
        except AttributeError:
            suffix = ""
        if suffix == ".teacher" and checkpoint_adapter.can_load(source):
            return checkpoint_adapter

    for key in (source.format, source.framework):
        normalized = _normalize(key)
        if normalized != "auto" and normalized in _ADAPTERS:
            return _ADAPTERS[normalized]

    for adapter in dict.fromkeys(_ADAPTERS.values()):
        if adapter.can_load(source):
            return adapter

    raise ValueError(
        "No teacher adapter can load the provided source. "
        f"framework={source.framework!r}, format={source.format!r}, path={source.path!r}. "
        f"Available adapters: {', '.join(available_teacher_adapters())}"
    )

load_teacher_state(source: TeacherSource, device: torch.device | str | None = None) -> TeacherState

Resolve an adapter and load a TeacherState from source.

Source code in recdistill/teachers/registry.py
def load_teacher_state(source: TeacherSource, device: torch.device | str | None = None) -> TeacherState:
    """Resolve an adapter and load a `TeacherState` from `source`."""
    adapter = resolve_teacher_adapter(source)
    return adapter.load(source, device=device)

Import Adapters

Import adapters convert external artifacts into the shared TeacherState format.

Adapter Accepted sources Output representation
CheckpointAdapter .teacher, .pt, .pth, .ckpt payloads. Embeddings, dense scores, or top-k scorer.
PredictionsJsonAdapter JSON prediction rows or column-oriented prediction exports. PrecomputedTopKScorer.
RecBolePthAdapter RecBole .pth checkpoints with embedding tensors. Embedding-backed TeacherState.

CheckpointAdapter

Load generic PyTorch checkpoints into TeacherState.

The adapter accepts serialized .teacher payloads, embedding dictionaries, dense score matrices, and top-k ranking payloads stored in .pt, .pth, .ckpt, or .teacher files.

Source code in recdistill/teachers/adapters/checkpoint.py
class CheckpointAdapter:
    """Load generic PyTorch checkpoints into `TeacherState`.

    The adapter accepts serialized `.teacher` payloads, embedding dictionaries,
    dense score matrices, and top-k ranking payloads stored in `.pt`, `.pth`,
    `.ckpt`, or `.teacher` files.
    """

    name = "checkpoint"

    def can_load(self, source: TeacherSource) -> bool:
        """Return `True` when the checkpoint payload can form a teacher state."""
        if _matches(source.format) or _matches(source.framework):
            return True
        if source.path is None:
            return False
        if Path(source.path).suffix.lower() not in {".teacher", ".pt", ".pth", ".ckpt"}:
            return False
        try:
            payload = _load_checkpoint(source.path)
        except Exception:
            return False
        return _can_build_teacher_state(payload)

    def load(self, source: TeacherSource, device: torch.device | str | None = None) -> TeacherState:
        """Load the checkpoint and convert it to `TeacherState`."""
        if source.path is None:
            raise ValueError("CheckpointAdapter requires --input.")
        payload = _load_checkpoint(source.path)
        state = _teacher_state_from_checkpoint(payload)
        state.metadata.update(source.metadata)
        state.metadata.setdefault("source_path", str(source.path))
        state.metadata.setdefault("source_format", "checkpoint")
        if source.model_name:
            state.metadata.setdefault("model_name", source.model_name)
        if device is not None:
            return state.to(device)
        return state

name = 'checkpoint' class-attribute instance-attribute

can_load(source: TeacherSource) -> bool

Return True when the checkpoint payload can form a teacher state.

Source code in recdistill/teachers/adapters/checkpoint.py
def can_load(self, source: TeacherSource) -> bool:
    """Return `True` when the checkpoint payload can form a teacher state."""
    if _matches(source.format) or _matches(source.framework):
        return True
    if source.path is None:
        return False
    if Path(source.path).suffix.lower() not in {".teacher", ".pt", ".pth", ".ckpt"}:
        return False
    try:
        payload = _load_checkpoint(source.path)
    except Exception:
        return False
    return _can_build_teacher_state(payload)

load(source: TeacherSource, device: torch.device | str | None = None) -> TeacherState

Load the checkpoint and convert it to TeacherState.

Source code in recdistill/teachers/adapters/checkpoint.py
def load(self, source: TeacherSource, device: torch.device | str | None = None) -> TeacherState:
    """Load the checkpoint and convert it to `TeacherState`."""
    if source.path is None:
        raise ValueError("CheckpointAdapter requires --input.")
    payload = _load_checkpoint(source.path)
    state = _teacher_state_from_checkpoint(payload)
    state.metadata.update(source.metadata)
    state.metadata.setdefault("source_path", str(source.path))
    state.metadata.setdefault("source_format", "checkpoint")
    if source.model_name:
        state.metadata.setdefault("model_name", source.model_name)
    if device is not None:
        return state.to(device)
    return state

PredictionsJsonAdapter

Import a teacher from JSON, TSV, or CSV prediction rows.

The prediction payload can be JSON (list of rows or dict with predictions), or TSV/CSV prediction exports containing user, item, and score/rating. Rows are converted into a PrecomputedTopKScorer.

Source code in recdistill/teachers/adapters/predictions_json.py
class PredictionsJsonAdapter:
    """Import a teacher from JSON, TSV, or CSV prediction rows.

    The prediction payload can be JSON (list of rows or dict with predictions),
    or TSV/CSV prediction exports containing user, item, and score/rating.
    Rows are converted into a `PrecomputedTopKScorer`.
    """

    name = "predictions_json"

    def can_load(self, source: TeacherSource) -> bool:
        """Return `True` for JSON/TSV/CSV prediction sources or matching format hints."""
        if _matches(source.format) or _matches(source.framework):
            return True
        if source.path is not None:
            suffix = Path(source.path).suffix.lower()
            return suffix in {".json", ".tsv", ".csv", ".txt"}
        return False

    def load(self, source: TeacherSource, device: torch.device | str | None = None) -> TeacherState:
        """Read the prediction payload and return a scorer-backed `TeacherState`."""
        if source.path is None:
            raise ValueError("PredictionsJsonAdapter requires --input.")
        path = Path(source.path)
        suffix = path.suffix.lower()
        if suffix in {".tsv", ".csv", ".txt"} or source.format in {"tsv", "csv", "predictions_tsv"}:
            rows = _read_tsv_rows(path)
        else:
            payload = json.loads(path.read_text(encoding="utf-8"))
            rows = _prediction_rows(payload)
        scorer, metadata = _topk_scorer_from_rows(rows, source.metadata)
        metadata.update(source.metadata)
        metadata.setdefault("source_path", str(source.path))
        metadata.setdefault("source_format", "predictions_json")
        if source.model_name:
            metadata.setdefault("model_name", source.model_name)
        state = TeacherState(metadata=metadata, scorer=scorer)
        if device is not None:
            return state.to(device)
        return state

name = 'predictions_json' class-attribute instance-attribute

can_load(source: TeacherSource) -> bool

Return True for JSON/TSV/CSV prediction sources or matching format hints.

Source code in recdistill/teachers/adapters/predictions_json.py
def can_load(self, source: TeacherSource) -> bool:
    """Return `True` for JSON/TSV/CSV prediction sources or matching format hints."""
    if _matches(source.format) or _matches(source.framework):
        return True
    if source.path is not None:
        suffix = Path(source.path).suffix.lower()
        return suffix in {".json", ".tsv", ".csv", ".txt"}
    return False

load(source: TeacherSource, device: torch.device | str | None = None) -> TeacherState

Read the prediction payload and return a scorer-backed TeacherState.

Source code in recdistill/teachers/adapters/predictions_json.py
def load(self, source: TeacherSource, device: torch.device | str | None = None) -> TeacherState:
    """Read the prediction payload and return a scorer-backed `TeacherState`."""
    if source.path is None:
        raise ValueError("PredictionsJsonAdapter requires --input.")
    path = Path(source.path)
    suffix = path.suffix.lower()
    if suffix in {".tsv", ".csv", ".txt"} or source.format in {"tsv", "csv", "predictions_tsv"}:
        rows = _read_tsv_rows(path)
    else:
        payload = json.loads(path.read_text(encoding="utf-8"))
        rows = _prediction_rows(payload)
    scorer, metadata = _topk_scorer_from_rows(rows, source.metadata)
    metadata.update(source.metadata)
    metadata.setdefault("source_path", str(source.path))
    metadata.setdefault("source_format", "predictions_json")
    if source.model_name:
        metadata.setdefault("model_name", source.model_name)
    state = TeacherState(metadata=metadata, scorer=scorer)
    if device is not None:
        return state.to(device)
    return state

RecBolePthAdapter

Import RecBole .pth checkpoints that contain embedding tensors.

The adapter inspects the state dict, finds compatible user/item embedding matrices, and stores the source tensor keys in teacher metadata.

Source code in recdistill/teachers/adapters/recbole_pth.py
class RecBolePthAdapter:
    """Import RecBole `.pth` checkpoints that contain embedding tensors.

    The adapter inspects the state dict, finds compatible user/item embedding
    matrices, and stores the source tensor keys in teacher metadata.
    """

    name = "recbole_pth"

    def can_load(self, source: TeacherSource) -> bool:
        """Return `True` when a `.pth` checkpoint exposes teacher embeddings."""
        if _matches(source.format) or _matches(source.framework):
            return True
        if source.path is None or Path(source.path).suffix.lower() != ".pth":
            return False
        try:
            state_dict = _extract_state_dict(_load_checkpoint(source.path))
        except Exception:
            return False
        return _find_user_item_embeddings(state_dict) is not None

    def load(self, source: TeacherSource, device: torch.device | str | None = None) -> TeacherState:
        """Load a RecBole checkpoint and convert its embeddings to `TeacherState`."""
        if source.path is None:
            raise ValueError("RecBolePthAdapter requires --input.")
        payload = _load_checkpoint(source.path)
        state_dict = _extract_state_dict(payload)
        tensors = _find_user_item_embeddings(state_dict)
        if tensors is None:
            raise ValueError("Unable to find compatible user/item embedding tensors in .pth checkpoint.")
        user_embeddings, item_embeddings, user_key, item_key = tensors
        metadata = {
            "representation": "embeddings",
            "source_path": str(source.path),
            "source_format": "recbole_pth",
            "user_embedding_key": user_key,
            "item_embedding_key": item_key,
        }
        metadata.update(source.metadata)
        if source.model_name:
            metadata.setdefault("model_name", source.model_name)
        state = TeacherState(
            user_embeddings=user_embeddings.detach().cpu().to(dtype=torch.float32),
            item_embeddings=item_embeddings.detach().cpu().to(dtype=torch.float32),
            metadata=metadata,
        )
        if device is not None:
            return state.to(device)
        return state

name = 'recbole_pth' class-attribute instance-attribute

can_load(source: TeacherSource) -> bool

Return True when a .pth checkpoint exposes teacher embeddings.

Source code in recdistill/teachers/adapters/recbole_pth.py
def can_load(self, source: TeacherSource) -> bool:
    """Return `True` when a `.pth` checkpoint exposes teacher embeddings."""
    if _matches(source.format) or _matches(source.framework):
        return True
    if source.path is None or Path(source.path).suffix.lower() != ".pth":
        return False
    try:
        state_dict = _extract_state_dict(_load_checkpoint(source.path))
    except Exception:
        return False
    return _find_user_item_embeddings(state_dict) is not None

load(source: TeacherSource, device: torch.device | str | None = None) -> TeacherState

Load a RecBole checkpoint and convert its embeddings to TeacherState.

Source code in recdistill/teachers/adapters/recbole_pth.py
def load(self, source: TeacherSource, device: torch.device | str | None = None) -> TeacherState:
    """Load a RecBole checkpoint and convert its embeddings to `TeacherState`."""
    if source.path is None:
        raise ValueError("RecBolePthAdapter requires --input.")
    payload = _load_checkpoint(source.path)
    state_dict = _extract_state_dict(payload)
    tensors = _find_user_item_embeddings(state_dict)
    if tensors is None:
        raise ValueError("Unable to find compatible user/item embedding tensors in .pth checkpoint.")
    user_embeddings, item_embeddings, user_key, item_key = tensors
    metadata = {
        "representation": "embeddings",
        "source_path": str(source.path),
        "source_format": "recbole_pth",
        "user_embedding_key": user_key,
        "item_embedding_key": item_key,
    }
    metadata.update(source.metadata)
    if source.model_name:
        metadata.setdefault("model_name", source.model_name)
    state = TeacherState(
        user_embeddings=user_embeddings.detach().cpu().to(dtype=torch.float32),
        item_embeddings=item_embeddings.detach().cpu().to(dtype=torch.float32),
        metadata=metadata,
    )
    if device is not None:
        return state.to(device)
    return state