torch-rechub 0.0.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

torch_rechub/serving/faiss.py

@@ -0,0 +1,154 @@
+"""FAISS-based vector index implementation for the retrieval stage."""
+
+import contextlib
+import typing as ty
+
+import faiss
+import torch
+
+from torch_rechub.types import FilePath
+
+from .base import BaseBuilder, BaseIndexer
+
+# Type for indexing methods.
+_FaissIndexType = ty.Literal["Flat", "HNSW", "IVF"]
+
+# Type for indexing metrics.
+_FaissMetric = ty.Literal["IP", "L2"]
+
+# Default indexing method.
+_DEFAULT_FAISS_INDEX_TYPE: _FaissIndexType = "Flat"
+
+# Default indexing metric.
+_DEFAULT_FAISS_METRIC: _FaissMetric = "L2"
+
+# Default number of clusters to build an IVF index.
+_DEFAULT_N_LISTS = 100
+
+# Default max number of neighbors to build an HNSW index.
+_DEFAULT_M = 32
+
+
+class FaissBuilder(BaseBuilder):
+    """Implement ``BaseBuilder`` for FAISS vector index construction."""
+
+    def __init__(
+        self,
+        index_type: _FaissIndexType = _DEFAULT_FAISS_INDEX_TYPE,
+        metric: _FaissMetric = _DEFAULT_FAISS_METRIC,
+        *,
+        m: int = _DEFAULT_M,
+        nlists: int = _DEFAULT_N_LISTS,
+        efSearch: ty.Optional[int] = None,
+        nprobe: ty.Optional[int] = None,
+    ) -> None:
+        """
+        Initialize a FAISS builder.
+
+        Parameters
+        ----------
+        index_type : ``"Flat"``, ``"HNSW"``, or ``"IVF"``, optional
+            The indexing index_type. Default to ``"Flat"``.
+        metric : ``"IP"``, ``"L2"``, optional
+            The indexing metric. Default to ``"L2"``.
+        m : int, optional
+            Max number of neighbors to build an HNSW index.
+        nlists : int, optional
+            Number of clusters to build an IVF index.
+        efSearch : int or None, optional
+            Number of candidate nodes during an HNSW search.
+        nprobe : int or None, optional
+            Number of clusters during an IVF search.
+        """
+        self._index_type_dsl = _build_index_type_dsl(index_type, m=m, nlists=nlists)
+        self._metric = _resolve_metric_type(metric)
+
+        self._efSearch = efSearch
+        self._nprobe = nprobe
+
+    @contextlib.contextmanager
+    def from_embeddings(
+        self,
+        embeddings: torch.Tensor,
+    ) -> ty.Generator["FaissIndexer",
+                      None,
+                      None]:
+        """Adhere to ``BaseBuilder.from_embeddings``."""
+        index: faiss.Index = faiss.index_factory(
+            embeddings.shape[1],
+            self._index_type_dsl,
+            self._metric,
+        )
+
+        if isinstance(index, faiss.IndexHNSW) and self._efSearch is not None:
+            index.hnsw.efSearch = self._efSearch
+
+        if isinstance(index, faiss.IndexIVF) and self._nprobe is not None:
+            index.nprobe = self._nprobe
+
+        index.train(embeddings)
+        index.add(embeddings)
+
+        try:
+            yield FaissIndexer(index)
+        finally:
+            pass
+
+    @contextlib.contextmanager
+    def from_index_file(
+        self,
+        index_file: FilePath,
+    ) -> ty.Generator["FaissIndexer",
+                      None,
+                      None]:
+        """Adhere to ``BaseBuilder.from_index_file``."""
+        index = faiss.read_index(str(index_file))
+
+        try:
+            yield FaissIndexer(index)
+        finally:
+            pass
+
+
+class FaissIndexer(BaseIndexer):
+    """FAISS-based implementation of ``BaseIndexer``."""
+
+    def __init__(self, index: faiss.Index) -> None:
+        """Initialize a FAISS indexer."""
+        self._index = index
+
+    def query(
+        self,
+        embeddings: torch.Tensor,
+        top_k: int,
+    ) -> tuple[torch.Tensor,
+               torch.Tensor]:
+        """Adhere to ``BaseIndexer.query``."""
+        dists, ids = self._index.search(embeddings.cpu().numpy(), top_k)
+        return torch.from_numpy(ids), torch.from_numpy(dists)
+
+    def save(self, file_path: FilePath) -> None:
+        """Adhere to ``BaseIndexer.save``."""
+        faiss.write_index(self._index, str(file_path))
+
+
+# helper functions
+
+
+def _build_index_type_dsl(index_type: _FaissIndexType, *, m: int, nlists: int) -> str:
+    """Build the index_type DSL passed to ``faiss.index_factory``."""
+    if index_type == "HNSW":
+        return f"{index_type}{m},Flat"
+
+    if index_type == "IVF":
+        return f"{index_type}{nlists},Flat"
+
+    return "Flat"
+
+
+def _resolve_metric_type(metric: _FaissMetric) -> int:
+    """Resolve the metric type from a string literal to an integer."""
+    if metric == "L2":
+        return ty.cast(int, faiss.METRIC_L2)
+
+    return ty.cast(int, faiss.METRIC_INNER_PRODUCT)

torch_rechub/serving/milvus.py

@@ -0,0 +1,215 @@
+"""Milvus-based vector index implementation for the retrieval stage."""
+
+import contextlib
+import typing as ty
+import uuid
+
+import numpy as np
+import pymilvus as milvus
+import torch
+
+from torch_rechub.types import FilePath
+
+from .base import BaseBuilder, BaseIndexer
+
+# Type for indexing methods.
+_MilvusIndexType = ty.Literal["FLAT", "HNSW", "IVF_FLAT"]
+
+# Type for indexing metrics.
+_MilvusMetric = ty.Literal["COSINE", "IP", "L2"]
+
+# Default indexing method.
+_DEFAULT_MILVUS_INDEX_TYPE: _MilvusIndexType = "FLAT"
+
+# Default indexing metric.
+_DEFAULT_MILVUS_METRIC: _MilvusMetric = "COSINE"
+
+# Default number of clusters to build an IVF index.
+_DEFAULT_N_LIST = 128
+
+# Default max number of neighbors to build an HNSW index.
+_DEFAULT_M = 30
+
+# Default name of Milvus database connection.
+_DEFAULT_NAME = "rechub"
+
+# Default host of Milvus instance.
+_DEFAULT_HOST = "localhost"
+
+# Default port of Milvus instance.
+_DEFAULT_PORT = 19530
+
+# Name of the embedding column in the Milvus database table.
+_EMBEDDING_COLUMN = "embedding"
+
+
+class MilvusBuilder(BaseBuilder):
+    """Implement ``BaseBuilder`` for Milvus vector index construction."""
+
+    def __init__(
+        self,
+        d: int,
+        index_type: _MilvusIndexType = _DEFAULT_MILVUS_INDEX_TYPE,
+        metric: _MilvusMetric = _DEFAULT_MILVUS_METRIC,
+        *,
+        m: int = _DEFAULT_M,
+        nlist: int = _DEFAULT_N_LIST,
+        ef: ty.Optional[int] = None,
+        nprobe: ty.Optional[int] = None,
+        name: str = _DEFAULT_NAME,
+        host: str = _DEFAULT_HOST,
+        port: int = _DEFAULT_PORT,
+    ) -> None:
+        """
+        Initialize a Milvus builder.
+
+        Parameters
+        ----------
+        d : int
+            The dimension of embeddings.
+        index_type : ``"FLAT"``, ``"HNSW"``, or ``"IVF_FLAT"``, optional
+            The indexing index_type. Default to ``"FLAT"``.
+        metric : ``"COSINE"``, ``"IP"``, or ``"L2"``, optional
+            The indexing metric. Default to ``"COSINE"``.
+        m : int, optional
+            Max number of neighbors to build an HNSW index.
+        nlist : int, optional
+            Number of clusters to build an IVF index.
+        ef : int or None, optional
+            Number of candidate nodes during an HNSW search.
+        nprobe : int or None, optional
+            Number of clusters during an IVF search.
+        name : str, optional
+            The name of connection. Each name corresponds to one connection.
+        host : str, optional
+            The host of Milvus instance. Default at "localhost".
+        port : int, optional
+            The port of Milvus instance. Default at 19530
+        """
+        self._d = d
+
+        # connection parameters
+        self._name = name
+        self._host = host
+        self._port = port
+
+        bparams: dict[str, ty.Any] = {}
+        qparams: dict[str, ty.Any] = {}
+
+        if index_type == "HNSW":
+            bparams.update(M=m)
+            if ef is not None:
+                qparams.update(ef=ef)
+
+        if index_type == "IVF_FLAT":
+            bparams.update(nlist=nlist)
+            if nprobe is not None:
+                qparams.update(nprobe=nprobe)
+
+        self._build_params = dict(
+            index_type=index_type,
+            metric_type=metric,
+            params=bparams,
+        )
+
+        self._query_params = dict(
+            metric_type=metric,
+            params=qparams,
+        )
+
+    @contextlib.contextmanager
+    def from_embeddings(
+        self,
+        embeddings: torch.Tensor,
+    ) -> ty.Generator["MilvusIndexer",
+                      None,
+                      None]:
+        """Adhere to ``BaseBuilder.from_embeddings``."""
+        milvus.connections.connect(self._name, host=self._host, port=self._port)
+        collection = self._build_collection(embeddings)
+
+        try:
+            yield MilvusIndexer(collection, self._query_params)
+        finally:
+            collection.drop()
+            milvus.connections.disconnect(self._name)
+
+    @contextlib.contextmanager
+    def from_index_file(
+        self,
+        index_file: FilePath,
+    ) -> ty.Generator["MilvusIndexer",
+                      None,
+                      None]:
+        """Adhere to ``BaseBuilder.from_index_file``."""
+        raise NotImplementedError("Milvus does not support index files!")
+
+    def _build_collection(self, embeddings: torch.Tensor) -> milvus.Collection:
+        """Build a Milvus collection with the current connection."""
+        fields = [
+            milvus.FieldSchema(
+                name="id",
+                dtype=milvus.DataType.INT64,
+                is_primary=True,
+            ),
+            milvus.FieldSchema(
+                name=_EMBEDDING_COLUMN,
+                dtype=milvus.DataType.FLOAT_VECTOR,
+                dim=self._d,
+            ),
+        ]
+
+        collection = milvus.Collection(
+            name=f"{self._name}_{uuid.uuid4().hex}",
+            schema=milvus.CollectionSchema(fields=fields),
+            using=self._name,
+        )
+
+        n, _ = embeddings.shape
+        collection.insert([np.arange(n, dtype=np.int64), embeddings.cpu().numpy()])
+        collection.create_index(_EMBEDDING_COLUMN, index_params=self._build_params)
+        collection.load()
+
+        return collection
+
+
+class MilvusIndexer(BaseIndexer):
+    """Milvus-based implementation of ``BaseIndexer``."""
+
+    def __init__(
+        self,
+        collection: milvus.Collection,
+        query_params: dict[str,
+                           ty.Any],
+    ) -> None:
+        """Initialize a Milvus indexer."""
+        self._collection = collection
+        self._query_params = query_params
+
+    def query(
+        self,
+        embeddings: torch.Tensor,
+        top_k: int,
+    ) -> tuple[torch.Tensor,
+               torch.Tensor]:
+        """Adhere to ``BaseIndexer.query``."""
+        results = self._collection.search(
+            data=embeddings.cpu().numpy(),
+            anns_field=_EMBEDDING_COLUMN,
+            param=self._query_params,
+            limit=top_k,
+        )
+
+        n, _ = embeddings.shape
+        nn_ids = np.zeros((n, top_k), dtype=np.int64)
+        nn_distances = np.zeros((n, top_k), dtype=np.float32)
+
+        for i, result in enumerate(results):
+            nn_ids[i] = result.ids
+            nn_distances[i] = result.distances
+
+        return torch.from_numpy(nn_ids), torch.from_numpy(nn_distances)
+
+    def save(self, file_path: FilePath) -> None:
+        """Adhere to ``BaseIndexer.save``."""
+        raise NotImplementedError("Milvus does not support index files!")

torch_rechub/trainers/ctr_trainer.py

@@ -185,7 +185,7 @@ class CTRTrainer(object):
                 predicts.extend(y_pred.tolist())
         return predicts
 
-    def export_onnx(self, output_path, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False):
+    def export_onnx(self, output_path, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False, onnx_export_kwargs=None):
         """Export the trained model to ONNX format.
 
         This method exports the ranking model (e.g., DeepFM, WideDeep, DCN) to ONNX format
@@ -202,6 +202,7 @@ class CTRTrainer(object):
             device (str, optional): Device for export ('cpu', 'cuda', etc.).
                 If None, defaults to 'cpu' for maximum compatibility.
             verbose (bool): Print export details (default: False).
+            onnx_export_kwargs (dict, optional): Extra kwargs forwarded to ``torch.onnx.export``.
 
         Returns:
             bool: True if export succeeded, False otherwise.
@@ -227,7 +228,16 @@ class CTRTrainer(object):
         export_device = device if device is not None else 'cpu'
 
         exporter = ONNXExporter(model, device=export_device)
-        return exporter.export(output_path=output_path, dummy_input=dummy_input, batch_size=batch_size, seq_length=seq_length, opset_version=opset_version, dynamic_batch=dynamic_batch, verbose=verbose)
+        return exporter.export(
+            output_path=output_path,
+            dummy_input=dummy_input,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            opset_version=opset_version,
+            dynamic_batch=dynamic_batch,
+            verbose=verbose,
+            onnx_export_kwargs=onnx_export_kwargs,
+        )
 
     def visualization(self, input_data=None, batch_size=2, seq_length=10, depth=3, show_shapes=True, expand_nested=True, save_path=None, graph_name="model", device=None, dpi=300, **kwargs):
         """Visualize the model's computation graph.

torch_rechub/trainers/match_trainer.py

@@ -215,7 +215,7 @@ class MatchTrainer(object):
                 predicts.append(y_pred.data)
         return torch.cat(predicts, dim=0)
 
-    def export_onnx(self, output_path, mode=None, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False):
+    def export_onnx(self, output_path, mode=None, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False, onnx_export_kwargs=None):
         """Export the trained matching model to ONNX format.
 
         This method exports matching/retrieval models (e.g., DSSM, YoutubeDNN, MIND)
@@ -237,6 +237,7 @@ class MatchTrainer(object):
             device (str, optional): Device for export ('cpu', 'cuda', etc.).
                 If None, defaults to 'cpu' for maximum compatibility.
             verbose (bool): Print export details (default: False).
+            onnx_export_kwargs (dict, optional): Extra kwargs forwarded to ``torch.onnx.export``.
 
         Returns:
             bool: True if export succeeded, False otherwise.
@@ -270,7 +271,17 @@ class MatchTrainer(object):
 
         try:
             exporter = ONNXExporter(model, device=export_device)
-            return exporter.export(output_path=output_path, mode=mode, dummy_input=dummy_input, batch_size=batch_size, seq_length=seq_length, opset_version=opset_version, dynamic_batch=dynamic_batch, verbose=verbose)
+            return exporter.export(
+                output_path=output_path,
+                mode=mode,
+                dummy_input=dummy_input,
+                batch_size=batch_size,
+                seq_length=seq_length,
+                opset_version=opset_version,
+                dynamic_batch=dynamic_batch,
+                verbose=verbose,
+                onnx_export_kwargs=onnx_export_kwargs,
+            )
         finally:
             # Restore original mode
             if hasattr(model, 'mode'):

torch_rechub/trainers/mtl_trainer.py

@@ -261,7 +261,7 @@ class MTLTrainer(object):
                 predicts.extend(y_preds.tolist())
         return predicts
 
-    def export_onnx(self, output_path, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False):
+    def export_onnx(self, output_path, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False, onnx_export_kwargs=None):
         """Export the trained multi-task model to ONNX format.
 
         This method exports multi-task learning models (e.g., MMOE, PLE, ESMM, SharedBottom)
@@ -283,6 +283,7 @@ class MTLTrainer(object):
             device (str, optional): Device for export ('cpu', 'cuda', etc.).
                 If None, defaults to 'cpu' for maximum compatibility.
             verbose (bool): Print export details (default: False).
+            onnx_export_kwargs (dict, optional): Extra kwargs forwarded to ``torch.onnx.export``.
 
         Returns:
             bool: True if export succeeded, False otherwise.
@@ -304,7 +305,16 @@ class MTLTrainer(object):
         export_device = device if device is not None else 'cpu'
 
         exporter = ONNXExporter(model, device=export_device)
-        return exporter.export(output_path=output_path, dummy_input=dummy_input, batch_size=batch_size, seq_length=seq_length, opset_version=opset_version, dynamic_batch=dynamic_batch, verbose=verbose)
+        return exporter.export(
+            output_path=output_path,
+            dummy_input=dummy_input,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            opset_version=opset_version,
+            dynamic_batch=dynamic_batch,
+            verbose=verbose,
+            onnx_export_kwargs=onnx_export_kwargs,
+        )
 
     def visualization(self, input_data=None, batch_size=2, seq_length=10, depth=3, show_shapes=True, expand_nested=True, save_path=None, graph_name="model", device=None, dpi=300, **kwargs):
         """Visualize the model's computation graph.

torch_rechub/trainers/seq_trainer.py

@@ -255,7 +255,7 @@ class SeqTrainer(object):
 
         return avg_loss, accuracy
 
-    def export_onnx(self, output_path, batch_size=2, seq_length=50, vocab_size=None, opset_version=14, dynamic_batch=True, device=None, verbose=False):
+    def export_onnx(self, output_path, batch_size=2, seq_length=50, vocab_size=None, opset_version=14, dynamic_batch=True, device=None, verbose=False, onnx_export_kwargs=None):
         """Export the trained sequence generation model to ONNX format.
 
         This method exports sequence generation models (e.g., HSTU) to ONNX format.
@@ -273,6 +273,7 @@ class SeqTrainer(object):
             device (str, optional): Device for export ('cpu', 'cuda', etc.).
                 If None, defaults to 'cpu' for maximum compatibility.
             verbose (bool): Print export details (default: False).
+            onnx_export_kwargs (dict, optional): Extra kwargs forwarded to ``torch.onnx.export``.
 
         Returns:
             bool: True if export succeeded, False otherwise.
@@ -321,20 +322,38 @@ class SeqTrainer(object):
 
         try:
             with torch.no_grad():
-                torch.onnx.export(
-                    model,
-                    (dummy_seq_tokens,
-                     dummy_seq_time_diffs),
-                    output_path,
-                    input_names=["seq_tokens",
-                                 "seq_time_diffs"],
-                    output_names=["output"],
-                    dynamic_axes=dynamic_axes,
-                    opset_version=opset_version,
-                    do_constant_folding=True,
-                    verbose=verbose,
-                    dynamo=False  # Use legacy exporter for dynamic_axes support
-                )
+                import inspect
+
+                export_kwargs = {
+                    "f": output_path,
+                    "input_names": ["seq_tokens",
+                                    "seq_time_diffs"],
+                    "output_names": ["output"],
+                    "dynamic_axes": dynamic_axes,
+                    "opset_version": opset_version,
+                    "do_constant_folding": True,
+                    "verbose": verbose,
+                }
+
+                if onnx_export_kwargs:
+                    overlap = set(export_kwargs.keys()) & set(onnx_export_kwargs.keys())
+                    overlap.discard("dynamo")
+                    if overlap:
+                        raise ValueError("onnx_export_kwargs contains keys that overlap with explicit args: "
+                                         f"{sorted(overlap)}. Please set them via export_onnx() parameters instead.")
+                    export_kwargs.update(onnx_export_kwargs)
+
+                # Auto-pick exporter:
+                # - dynamic_axes present => prefer legacy exporter (dynamo=False) for dynamic batch/seq
+                # - otherwise prefer dynamo exporter (dynamo=True) on newer torch
+                sig = inspect.signature(torch.onnx.export)
+                if "dynamo" in sig.parameters:
+                    if "dynamo" not in export_kwargs:
+                        export_kwargs["dynamo"] = False if dynamic_axes is not None else True
+                else:
+                    export_kwargs.pop("dynamo", None)
+
+                torch.onnx.export(model, (dummy_seq_tokens, dummy_seq_time_diffs), **export_kwargs)
 
             if verbose:
                 print(f"Successfully exported ONNX model to: {output_path}")

torch_rechub/types.py

@@ -0,0 +1,5 @@
+import os
+import typing as ty
+
+#: Type for path to a file.
+FilePath = ty.Union[str, os.PathLike]