torch-rechub 0.0.5__py3-none-any.whl → 0.1.0__py3-none-any.whl

torch_rechub/basic/tracking.py

@@ -0,0 +1,198 @@
+"""Experiment tracking utilities for Torch-RecHub.
+
+This module exposes lightweight adapters for common visualization and
+experiment tracking tools, namely Weights & Biases (wandb), SwanLab, and
+TensorBoardX.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Union
+
+
+class BaseLogger(ABC):
+    """Base interface for experiment tracking backends.
+
+    Methods
+    -------
+    log_metrics(metrics, step=None)
+        Record scalar metrics at a given step.
+    log_hyperparams(params)
+        Store hyperparameters and run configuration.
+    finish()
+        Flush pending logs and release resources.
+    """
+
+    @abstractmethod
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        """Log metrics to the tracking backend.
+
+        Parameters
+        ----------
+        metrics : dict of str to Any
+            Metric name-value pairs to record.
+        step : int, optional
+            Explicit global step or epoch index. When ``None``, the backend
+            uses its own default step handling.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        """Log experiment hyperparameters.
+
+        Parameters
+        ----------
+        params : dict of str to Any
+            Hyperparameters or configuration values to persist with the run.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def finish(self) -> None:
+        """Finalize logging and free any backend resources."""
+        raise NotImplementedError
+
+
+class WandbLogger(BaseLogger):
+    """Weights & Biases logger implementation.
+
+    Parameters
+    ----------
+    project : str
+        Name of the wandb project to log to.
+    name : str, optional
+        Display name for the run.
+    config : dict, optional
+        Initial hyperparameter configuration to record.
+    tags : list of str, optional
+        Optional tags for grouping runs.
+    notes : str, optional
+        Long-form notes shown in the run overview.
+    dir : str, optional
+        Local directory for wandb artifacts and cache.
+    **kwargs : dict
+        Additional keyword arguments forwarded to ``wandb.init``.
+
+    Raises
+    ------
+    ImportError
+        If ``wandb`` is not installed in the current environment.
+    """
+
+    def __init__(self, project: str, name: Optional[str] = None, config: Optional[Dict[str, Any]] = None, tags: Optional[List[str]] = None, notes: Optional[str] = None, dir: Optional[str] = None, **kwargs):
+        try:
+            import wandb
+            self._wandb = wandb
+        except ImportError:
+            raise ImportError("wandb is not installed. Install it with: pip install wandb")
+
+        self.run = self._wandb.init(project=project, name=name, config=config, tags=tags, notes=notes, dir=dir, **kwargs)
+
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        if step is not None:
+            self._wandb.log(metrics, step=step)
+        else:
+            self._wandb.log(metrics)
+
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        if self.run is not None:
+            self.run.config.update(params)
+
+    def finish(self) -> None:
+        if self.run is not None:
+            self.run.finish()
+
+
+class SwanLabLogger(BaseLogger):
+    """SwanLab logger implementation.
+
+    Parameters
+    ----------
+    project : str, optional
+        Project identifier for grouping experiments.
+    experiment_name : str, optional
+        Display name for the experiment or run.
+    description : str, optional
+        Text description shown alongside the run.
+    config : dict, optional
+        Hyperparameters or configuration to log at startup.
+    logdir : str, optional
+        Directory where logs and artifacts are stored.
+    **kwargs : dict
+        Additional keyword arguments forwarded to ``swanlab.init``.
+
+    Raises
+    ------
+    ImportError
+        If ``swanlab`` is not installed in the current environment.
+    """
+
+    def __init__(self, project: Optional[str] = None, experiment_name: Optional[str] = None, description: Optional[str] = None, config: Optional[Dict[str, Any]] = None, logdir: Optional[str] = None, **kwargs):
+        try:
+            import swanlab
+            self._swanlab = swanlab
+        except ImportError:
+            raise ImportError("swanlab is not installed. Install it with: pip install swanlab")
+
+        self.run = self._swanlab.init(project=project, experiment_name=experiment_name, description=description, config=config, logdir=logdir, **kwargs)
+
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        if step is not None:
+            self._swanlab.log(metrics, step=step)
+        else:
+            self._swanlab.log(metrics)
+
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        if self.run is not None:
+            self.run.config.update(params)
+
+    def finish(self) -> None:
+        self._swanlab.finish()
+
+
+class TensorBoardXLogger(BaseLogger):
+    """TensorBoardX logger implementation.
+
+    Parameters
+    ----------
+    log_dir : str
+        Directory where event files will be written.
+    comment : str, default=""
+        Comment appended to the log directory name.
+    **kwargs : dict
+        Additional keyword arguments forwarded to
+        ``tensorboardX.SummaryWriter``.
+
+    Raises
+    ------
+    ImportError
+        If ``tensorboardX`` is not installed in the current environment.
+    """
+
+    def __init__(self, log_dir: str, comment: str = "", **kwargs):
+        try:
+            from tensorboardX import SummaryWriter
+            self._SummaryWriter = SummaryWriter
+        except ImportError:
+            raise ImportError("tensorboardX is not installed. Install it with: pip install tensorboardX")
+
+        self.writer = self._SummaryWriter(log_dir=log_dir, comment=comment, **kwargs)
+        self._step = 0
+
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        if step is None:
+            step = self._step
+            self._step += 1
+
+        for key, value in metrics.items():
+            if value is not None:
+                if isinstance(value, (int, float)):
+                    self.writer.add_scalar(key, value, step)
+
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        hparam_str = "\n".join([f"{k}: {v}" for k, v in params.items()])
+        self.writer.add_text("hyperparameters", hparam_str, 0)
+
+    def finish(self) -> None:
+        if self.writer is not None:
+            self.writer.close()

torch_rechub/data/convert.py

@@ -0,0 +1,67 @@
+"""Utilities for converting array-like data structures into PyTorch tensors."""
+
+import numpy.typing as npt
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.types as pt
+import torch
+
+
+def pa_array_to_tensor(arr: pa.Array) -> torch.Tensor:
+    """
+    Convert a PyArrow array to a PyTorch tensor.
+
+    Parameters
+    ----------
+    arr : pa.Array
+        The given PyArrow array.
+
+    Returns
+    -------
+    torch.Tensor: The result PyTorch tensor.
+
+    Raises
+    ------
+    TypeError
+        if the array type or the value type (when nested) is unsupported.
+    ValueError
+        if the nested array is ragged (unequal lengths of each row).
+    """
+    if _is_supported_scalar(arr.type):
+        arr = pc.cast(arr, pa.float32())
+        return torch.from_numpy(_to_writable_numpy(arr))
+
+    if not _is_supported_list(arr.type):
+        raise TypeError(f"Unsupported array type: {arr.type}")
+
+    if not _is_supported_scalar(val_type := arr.type.value_type):
+        raise TypeError(f"Unsupported value type in the nested array: {val_type}")
+
+    if len(pc.unique(pc.list_value_length(arr))) > 1:
+        raise ValueError("Cannot convert the ragged nested array.")
+
+    arr = pc.cast(arr, pa.list_(pa.float32()))
+    np_arr = _to_writable_numpy(arr.values)  # type: ignore[attr-defined]
+
+    # For empty list-of-lists, define output shape as (0, 0); otherwise infer width.
+    return torch.from_numpy(np_arr.reshape(len(arr), -1 if len(arr) > 0 else 0))
+
+
+# helper functions
+
+
+def _is_supported_list(t: pa.DataType) -> bool:
+    """Check if the given PyArrow data type is a supported list."""
+    return pt.is_fixed_size_list(t) or pt.is_large_list(t) or pt.is_list(t)
+
+
+def _is_supported_scalar(t: pa.DataType) -> bool:
+    """Check if the given PyArrow data type is a supported scalar type."""
+    return pt.is_boolean(t) or pt.is_floating(t) or pt.is_integer(t) or pt.is_null(t)
+
+
+def _to_writable_numpy(arr: pa.Array) -> npt.NDArray:
+    """Dump a PyArrow array into a writable NumPy array."""
+    # Force the NumPy array to be writable. PyArrow's to_numpy() often returns a
+    # read-only view for zero-copy, which PyTorch's from_numpy() does not support.
+    return arr.to_numpy(writable=True, zero_copy_only=False)

torch_rechub/data/dataset.py

@@ -0,0 +1,107 @@
+"""Dataset implementations providing streaming, batch-wise data access for PyTorch."""
+
+import typing as ty
+
+import pyarrow.dataset as pd
+import torch
+from torch.utils.data import IterableDataset, get_worker_info
+
+from torch_rechub.types import FilePath
+
+from .convert import pa_array_to_tensor
+
+# The default batch size when reading a Parquet dataset
+_DEFAULT_BATCH_SIZE = 1024
+
+
+class ParquetIterableDataset(IterableDataset):
+    """Stream Parquet data as PyTorch tensors.
+
+    Parameters
+    ----------
+    file_paths : list[FilePath]
+        Paths to Parquet files.
+    columns : list[str], optional
+        Columns to select; if ``None``, read all columns.
+    batch_size : int, default _DEFAULT_BATCH_SIZE
+        Rows per streamed batch.
+
+    Notes
+    -----
+    Reads lazily; no full Parquet load. Each worker gets a partition, builds its
+    own PyArrow Dataset/Scanner, and yields dicts of column tensors batch by batch.
+
+    Examples
+    --------
+    >>> ds = ParquetIterableDataset(
+    ...     ["/data/train1.parquet", "/data/train2.parquet"],
+    ...     columns=["x", "y", "label"],
+    ...     batch_size=1024,
+    ... )
+    >>> loader = DataLoader(ds, batch_size=None)
+    >>> for batch in loader:
+    ...     x, y, label = batch["x"], batch["y"], batch["label"]
+    ...     ...
+    """
+
+    def __init__(
+        self,
+        file_paths: ty.Sequence[FilePath],
+        /,
+        columns: ty.Optional[ty.Sequence[str]] = None,
+        batch_size: int = _DEFAULT_BATCH_SIZE,
+    ) -> None:
+        """Initialize this instance."""
+        self._file_paths = tuple(map(str, file_paths))
+        self._columns = None if columns is None else tuple(columns)
+        self._batch_size = batch_size
+
+    def __iter__(self) -> ty.Iterator[dict[str, torch.Tensor]]:
+        """Stream Parquet data as mapped PyTorch tensors.
+
+        Builds a PyArrow Dataset from the current worker's file partition, then
+        lazily scans selected columns. Each batch becomes a dict of Torch tensors.
+
+        Returns
+        -------
+        Iterator[dict[str, torch.Tensor]]
+            One converted batch at a time.
+        """
+        if not (partition := self._get_partition()):
+            return
+
+        # Build the dataset for the current worker.
+        ds = pd.dataset(partition, format="parquet")
+
+        # Create a scanner. This does not read data.
+        columns = None if self._columns is None else list(self._columns)
+        scanner = ds.scanner(columns=columns, batch_size=self._batch_size)
+
+        for batch in scanner.to_batches():
+            data_dict: dict[str, torch.Tensor] = {}
+            for name, array in zip(batch.column_names, batch.columns):
+                data_dict[name] = pa_array_to_tensor(array)
+            yield data_dict
+
+    # private interfaces
+
+    def _get_partition(self) -> tuple[str, ...]:
+        """Get file partition for the current worker.
+
+        Splits file paths into contiguous partitions by number of workers and worker ID.
+        In the main process (no worker info), returns all paths.
+
+        Returns
+        -------
+        tuple[str, ...]
+            Partition of file paths for this worker.
+        """
+        if (info := get_worker_info()) is None:
+            return self._file_paths
+
+        n = len(self._file_paths)
+        per_worker = (n + info.num_workers - 1) // info.num_workers
+
+        start = info.id * per_worker
+        end = n if (end := start + per_worker) > n else end
+        return self._file_paths[start:end]

torch_rechub/models/generative/hstu.py

@@ -10,39 +10,54 @@ from torch_rechub.utils.hstu_utils import RelPosBias
 
 
 class HSTUModel(nn.Module):
-    """HSTU: Hierarchical Sequential Transduction Units model.
-
-    Autoregressive generative recommendation model for sequential data.
-    This module stacks multiple ``HSTUBlock`` layers to capture long-range
-    dependencies in user interaction sequences and predicts the next item.
-
-    Args:
-        vocab_size (int): Vocabulary size (number of distinct items, including PAD).
-        d_model (int): Hidden dimension of the model. Default: 512.
-        n_heads (int): Number of attention heads. Default: 8.
-        n_layers (int): Number of stacked HSTU layers. Default: 4.
-        dqk (int): Dimension of query/key vectors per head. Default: 64.
-        dv (int): Dimension of value vectors per head. Default: 64.
-        max_seq_len (int): Maximum sequence length. Default: 256.
-        dropout (float): Dropout rate applied in the model. Default: 0.1.
-        use_rel_pos_bias (bool): Whether to use relative position bias. Default: True.
-        use_time_embedding (bool): Whether to use time-difference embeddings. Default: True.
-        num_time_buckets (int): Number of time buckets for time embeddings. Default: 2048.
-        time_bucket_fn (str): Function used to bucketize time differences, ``"sqrt"``
-            or ``"log"``. Default: ``"sqrt"``.
-
-    Shape:
-        - Input: ``x`` of shape ``(batch_size, seq_len)``; optional ``time_diffs``
-          of shape ``(batch_size, seq_len)`` representing time differences in seconds.
-        - Output: Logits of shape ``(batch_size, seq_len, vocab_size)``.
-
-    Example:
-        >>> model = HSTUModel(vocab_size=100000, d_model=512)
-        >>> x = torch.randint(0, 100000, (32, 256))
-        >>> time_diffs = torch.randint(0, 86400, (32, 256))
-        >>> logits = model(x, time_diffs)
-        >>> logits.shape
-        torch.Size([32, 256, 100000])
+    """HSTU: Hierarchical Sequential Transduction Units.
+
+    Autoregressive generative recommender that stacks ``HSTUBlock`` layers to
+    capture long-range dependencies and predict the next item.
+
+    Parameters
+    ----------
+    vocab_size : int
+        Vocabulary size (items incl. PAD).
+    d_model : int, default=512
+        Hidden dimension.
+    n_heads : int, default=8
+        Attention heads.
+    n_layers : int, default=4
+        Number of stacked HSTU layers.
+    dqk : int, default=64
+        Query/key dim per head.
+    dv : int, default=64
+        Value dim per head.
+    max_seq_len : int, default=256
+        Maximum sequence length.
+    dropout : float, default=0.1
+        Dropout rate.
+    use_rel_pos_bias : bool, default=True
+        Use relative position bias.
+    use_time_embedding : bool, default=True
+        Use time-difference embeddings.
+    num_time_buckets : int, default=2048
+        Number of time buckets for time embeddings.
+    time_bucket_fn : {'sqrt', 'log'}, default='sqrt'
+        Bucketization function for time differences.
+
+    Shape
+    -----
+    Input
+        x : ``(batch_size, seq_len)``
+        time_diffs : ``(batch_size, seq_len)``, optional (seconds).
+    Output
+        logits : ``(batch_size, seq_len, vocab_size)``
+
+    Examples
+    --------
+    >>> model = HSTUModel(vocab_size=100000, d_model=512)
+    >>> x = torch.randint(0, 100000, (32, 256))
+    >>> time_diffs = torch.randint(0, 86400, (32, 256))
+    >>> logits = model(x, time_diffs)
+    >>> logits.shape
+    torch.Size([32, 256, 100000])
     """
 
     def __init__(self, vocab_size, d_model=512, n_heads=8, n_layers=4, dqk=64, dv=64, max_seq_len=256, dropout=0.1, use_rel_pos_bias=True, use_time_embedding=True, num_time_buckets=2048, time_bucket_fn='sqrt'):

torch_rechub/serving/__init__.py

@@ -0,0 +1,50 @@
+import typing as ty
+
+from .annoy import AnnoyBuilder
+from .base import BaseBuilder
+from .faiss import FaissBuilder
+from .milvus import MilvusBuilder
+
+# Type for supported retrieval models.
+_RetrievalModel = ty.Literal["annoy", "faiss", "milvus"]
+
+
+def builder_factory(model: _RetrievalModel, **builder_config) -> BaseBuilder:
+    """
+    Factory function for creating a vector index builder.
+
+    This function instantiates and returns a concrete implementation of ``BaseBuilder``
+    based on the specified retrieval backend. The returned builder is responsible for
+    constructing or loading the underlying ANN index via its own ``from_embeddings`` or
+    ``from_index_file`` method.
+
+    Parameters
+    ----------
+    model : "annoy", "faiss", or "milvus"
+        The retrieval backend to use.
+    **builder_config
+        Keyword arguments passed directly to the selected builder constructor.
+
+    Returns
+    -------
+    BaseBuilder
+        A concrete builder instance corresponding to the specified retrieval backend.
+
+    Raises
+    ------
+    NotImplementedError
+        if the specified retrieval model is not supported.
+    """
+    if model == "annoy":
+        return AnnoyBuilder(**builder_config)
+
+    if model == "faiss":
+        return FaissBuilder(**builder_config)
+
+    if model == "milvus":
+        return MilvusBuilder(**builder_config)
+
+    raise NotImplementedError(f"{model=} is not implemented yet!")
+
+
+__all__ = ["builder_factory"]

torch_rechub/serving/annoy.py

@@ -0,0 +1,133 @@
+"""ANNOY-based vector index implementation for the retrieval stage."""
+
+import contextlib
+import typing as ty
+
+import annoy
+import numpy as np
+import torch
+
+from torch_rechub.types import FilePath
+
+from .base import BaseBuilder, BaseIndexer
+
+# Type for distance metrics for the ANNOY index.
+_AnnoyMetric = ty.Literal["angular", "euclidean", "dot"]
+
+# Default distance metric used by ANNOY.
+_DEFAULT_METRIC: _AnnoyMetric = "angular"
+
+# Default number of trees to build in the ANNOY index.
+_DEFAULT_N_TREES = 10
+
+# Default number of worker threads for building the ANNOY index.
+_DEFAULT_THREADS = -1
+
+# Default number of nodes to inspect during an ANNOY search.
+_DEFAULT_SEARCHK = -1
+
+
+class AnnoyBuilder(BaseBuilder):
+    """ANNOY-based implementation of ``BaseBuilder``."""
+
+    def __init__(
+        self,
+        d: int,
+        metric: _AnnoyMetric = _DEFAULT_METRIC,
+        *,
+        n_trees: int = _DEFAULT_N_TREES,
+        threads: int = _DEFAULT_THREADS,
+        searchk: int = _DEFAULT_SEARCHK,
+    ) -> None:
+        """
+        Initialize a ANNOY builder.
+
+        Parameters
+        ----------
+        d : int
+            The dimension of embeddings.
+        metric : ``"angular"``, ``"euclidean"``, or ``"dot"``, optional
+            The indexing metric. Default to ``"angular"``.
+        n_trees : int, optional
+            Number of trees to build an ANNOY index.
+        threads : int, optional
+            Number of worker threads to build an ANNOY index.
+        searchk : int, optional
+            Number of nodes to inspect during an ANNOY search.
+        """
+        self._d = d
+        self._metric = metric
+
+        self._n_trees = n_trees
+        self._threads = threads
+        self._searchk = searchk
+
+    @contextlib.contextmanager
+    def from_embeddings(
+        self,
+        embeddings: torch.Tensor,
+    ) -> ty.Generator["AnnoyIndexer",
+                      None,
+                      None]:
+        """Adhere to ``BaseBuilder.from_embeddings``."""
+        index = annoy.AnnoyIndex(self._d, metric=self._metric)
+
+        for idx, emb in enumerate(embeddings):
+            index.add_item(idx, emb)
+
+        index.build(self._n_trees, n_jobs=self._threads)
+
+        try:
+            yield AnnoyIndexer(index, self._searchk)
+        finally:
+            index.unload()
+
+    @contextlib.contextmanager
+    def from_index_file(
+        self,
+        index_file: FilePath,
+    ) -> ty.Generator["AnnoyIndexer",
+                      None,
+                      None]:
+        """Adhere to ``BaseBuilder.from_index_file``."""
+        index = annoy.AnnoyIndex(self._d, metric=self._metric)
+        index.load(str(index_file))
+
+        try:
+            yield AnnoyIndexer(index, searchk=self._searchk)
+        finally:
+            index.unload()
+
+
+class AnnoyIndexer(BaseIndexer):
+    """ANNOY-based implementation of ``BaseIndexer``."""
+
+    def __init__(self, index: annoy.AnnoyIndex, searchk: int) -> None:
+        """Initialize a ANNOY indexer."""
+        self._index = index
+        self._searchk = searchk
+
+    def query(
+        self,
+        embeddings: torch.Tensor,
+        top_k: int,
+    ) -> tuple[torch.Tensor,
+               torch.Tensor]:
+        """Adhere to ``BaseIndexer.query``."""
+        n, _ = embeddings.shape
+        nn_ids = np.zeros((n, top_k), dtype=np.int64)
+        nn_distances = np.zeros((n, top_k), dtype=np.float32)
+
+        for idx, emb in enumerate(embeddings):
+            nn_ids[idx], nn_distances[idx] = self._index.get_nns_by_vector(
+                emb.cpu().numpy(),
+                top_k,
+                search_k=self._searchk,
+                include_distances=True,
+            )
+
+        return torch.from_numpy(nn_ids), torch.from_numpy(nn_distances)
+
+    def save(self, file_path: FilePath) -> None:
+        """Adhere to ``BaseIndexer.save``."""
+        self._index.save(str(file_path))