nextrec 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

nextrec/loss/__init__.py

@@ -37,6 +37,5 @@ __all__ = [
     # Utilities
     "get_loss_fn",
     "get_loss_kwargs",
-    "validate_training_mode",
     "VALID_TASK_TYPES",
 ]

nextrec/loss/loss_utils.py

@@ -21,138 +21,69 @@ from nextrec.loss.pointwise import (
     WeightedBCELoss,
 )
 
-# Valid task types for validation
+
 VALID_TASK_TYPES = [
-    "binary",
-    "multiclass",
-    "regression",
-    "multivariate_regression",
-    "match",
-    "ranking",
-    "multitask",
-    "multilabel",
+    "binary",      
+    "multiclass", 
+    "multilabel", 
+    "regression",  
 ]
 
+def _build_cb_focal(kw):
+    if "class_counts" not in kw:
+        raise ValueError("class_balanced_focal requires class_counts")
+    return ClassBalancedFocalLoss(**kw)
 
-def get_loss_fn(
-    task_type: str = "binary",
-    training_mode: str | None = None,
-    loss: str | nn.Module | None = None,
-    **loss_kwargs,
-) -> nn.Module:
-    """
-    Get loss function based on task type and training mode.
-    """
 
+def get_loss_fn(loss=None, **kw):
     if isinstance(loss, nn.Module):
         return loss
-
-    # Common mappings
-    if task_type == "match":
-        return _get_match_loss(training_mode, loss, **loss_kwargs)
-
-    if task_type in ["ranking", "multitask", "binary", "multilabel"]:
-        return _get_classification_loss(loss, **loss_kwargs)
-
-    if task_type == "multiclass":
-        return _get_multiclass_loss(loss, **loss_kwargs)
-
-    if task_type == "regression":
-        if loss is None or loss == "mse":
-            return nn.MSELoss(**loss_kwargs)
-        if loss == "mae":
-            return nn.L1Loss(**loss_kwargs)
-        if isinstance(loss, str):
-            raise ValueError(f"Unsupported regression loss: {loss}")
-
-    raise ValueError(f"Unsupported task_type: {task_type}")
-
-
-def _get_match_loss(training_mode: str | None, loss: str | None, **loss_kwargs) -> nn.Module:
-    if training_mode == "pointwise":
-        if loss is None or loss in {"bce", "binary_crossentropy"}:
-            return nn.BCELoss(**loss_kwargs)
-        if loss == "weighted_bce":
-            return WeightedBCELoss(**loss_kwargs)
-        if loss == "focal":
-            return FocalLoss(**loss_kwargs)
-        if loss == "class_balanced_focal":
-            return _build_cb_focal(loss_kwargs)
-        if loss == "cosine_contrastive":
-            return CosineContrastiveLoss(**loss_kwargs)
-        if isinstance(loss, str):
-            raise ValueError(f"Unsupported pointwise loss: {loss}")
-
-    if training_mode == "pairwise":
-        if loss is None or loss == "bpr":
-            return BPRLoss(**loss_kwargs)
-        if loss == "hinge":
-            return HingeLoss(**loss_kwargs)
-        if loss == "triplet":
-            return TripletLoss(**loss_kwargs)
-        if isinstance(loss, str):
-            raise ValueError(f"Unsupported pairwise loss: {loss}")
-
-    if training_mode == "listwise":
-        if loss is None or loss in {"sampled_softmax", "softmax"}:
-            return SampledSoftmaxLoss(**loss_kwargs)
-        if loss == "infonce":
-            return InfoNCELoss(**loss_kwargs)
-        if loss == "listnet":
-            return ListNetLoss(**loss_kwargs)
-        if loss == "listmle":
-            return ListMLELoss(**loss_kwargs)
-        if loss == "approx_ndcg":
-            return ApproxNDCGLoss(**loss_kwargs)
-        if loss in {"crossentropy", "ce"}:
-            return nn.CrossEntropyLoss(**loss_kwargs)
-        if isinstance(loss, str):
-            raise ValueError(f"Unsupported listwise loss: {loss}")
-
-    raise ValueError(f"Unknown training_mode: {training_mode}")
-
-
-def _get_classification_loss(loss: str | None, **loss_kwargs) -> nn.Module:
-    if loss is None or loss in {"bce", "binary_crossentropy"}:
-        return nn.BCELoss(**loss_kwargs)
+    if loss is None:
+        raise ValueError("loss must be provided explicitly")
+    if loss in ["bce", "binary_crossentropy"]:
+        return nn.BCELoss(**kw)
     if loss == "weighted_bce":
-        return WeightedBCELoss(**loss_kwargs)
-    if loss == "focal":
-        return FocalLoss(**loss_kwargs)
-    if loss == "class_balanced_focal":
-        return _build_cb_focal(loss_kwargs)
+        return WeightedBCELoss(**kw)
+    if loss in ["focal", "focal_loss"]:
+        return FocalLoss(**kw)
+    if loss in ["cb_focal", "class_balanced_focal"]:
+        return _build_cb_focal(kw)
+    if loss in ["crossentropy", "ce"]:
+        return nn.CrossEntropyLoss(**kw)
     if loss == "mse":
-        return nn.MSELoss(**loss_kwargs)
+        return nn.MSELoss(**kw)
     if loss == "mae":
-        return nn.L1Loss(**loss_kwargs)
-    if loss in {"crossentropy", "ce"}:
-        return nn.CrossEntropyLoss(**loss_kwargs)
-    if isinstance(loss, str):
-        raise ValueError(f"Unsupported loss function: {loss}")
-    raise ValueError("Loss must be specified for classification task.")
-
-
-def _get_multiclass_loss(loss: str | None, **loss_kwargs) -> nn.Module:
-    if loss is None or loss in {"crossentropy", "ce"}:
-        return nn.CrossEntropyLoss(**loss_kwargs)
-    if loss == "focal":
-        return FocalLoss(**loss_kwargs)
-    if loss == "class_balanced_focal":
-        return _build_cb_focal(loss_kwargs)
-    if isinstance(loss, str):
-        raise ValueError(f"Unsupported multiclass loss: {loss}")
-    raise ValueError("Loss must be specified for multiclass task.")
-
-
-def _build_cb_focal(loss_kwargs: dict) -> ClassBalancedFocalLoss:
-    if "class_counts" not in loss_kwargs:
-        raise ValueError("class_balanced_focal requires `class_counts` argument.")
-    return ClassBalancedFocalLoss(**loss_kwargs)
-
+        return nn.L1Loss(**kw)
+
+    # Pairwise ranking Loss
+    if loss == "bpr":
+        return BPRLoss(**kw)
+    if loss == "hinge":
+        return HingeLoss(**kw)
+    if loss == "triplet":
+        return TripletLoss(**kw)
+
+    # Listwise ranking Loss
+    if loss in ["sampled_softmax", "softmax"]:
+        return SampledSoftmaxLoss(**kw)
+    if loss == "infonce":
+        return InfoNCELoss(**kw)
+    if loss == "listnet":
+        return ListNetLoss(**kw)
+    if loss == "listmle":
+        return ListMLELoss(**kw)
+    if loss == "approx_ndcg":
+        return ApproxNDCGLoss(**kw)
+
+    raise ValueError(f"Unsupported loss: {loss}")
 
 def get_loss_kwargs(loss_params: dict | list[dict] | None, index: int = 0) -> dict:
     """
-    Resolve per-task loss kwargs from a dict or list of dicts.
+    解析每个 head 对应的 loss_kwargs。
+
+    - loss_params 为 None       -> {}
+    - loss_params 为 dict       -> 所有 head 共用
+    - loss_params 为 list[dict] -> 用 loss_params[index]（若存在且非 None），否则 {}
     """
     if loss_params is None:
         return {}
@@ -160,4 +91,4 @@ def get_loss_kwargs(loss_params: dict | list[dict] | None, index: int = 0) -> di
         if index < len(loss_params) and loss_params[index] is not None:
             return loss_params[index]
         return {}
-    return loss_params
+    return loss_params

nextrec/models/multi_task/esmm.py

@@ -40,7 +40,7 @@ class ESMM(BaseModel):
                  ctr_params: dict,
                  cvr_params: dict,
                  target: list[str] = ['ctr', 'ctcvr'],  # Note: ctcvr = ctr * cvr
-                 task: str | list[str] = 'binary',
+                 task: list[str] = ['binary', 'binary'],
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[str | nn.Module] | None = "bce",

nextrec/models/ranking/masknet.py

@@ -144,7 +144,7 @@ class MaskNet(BaseModel):
 
     @property
     def task_type(self):
-        return "binary"
+        return "binary_classification"
 
     def __init__(
         self,

nextrec/utils/__init__.py

@@ -1,14 +1,17 @@
 from .optimizer import get_optimizer, get_scheduler
 from .initializer import get_initializer
 from .embedding import get_auto_embedding_dim
-from . import optimizer, initializer, embedding
+from .common import resolve_device
+from . import optimizer, initializer, embedding, common
 
 __all__ = [
     'get_optimizer',
     'get_scheduler',
     'get_initializer',
     'get_auto_embedding_dim',
+    'resolve_device',
     'optimizer',
     'initializer',
     'embedding',
+    'common',
 ]

nextrec/utils/common.py

@@ -0,0 +1,16 @@
+import torch
+import platform
+
+def resolve_device() -> str:
+    """Select a usable device with graceful fallback."""
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        mac_ver = platform.mac_ver()[0]
+        try:
+            major, minor = (int(x) for x in mac_ver.split(".")[:2])
+        except Exception:
+            major, minor = 0, 0
+        if major >= 14:
+            return "mps"
+    return "cpu"

{nextrec-0.2.4.dist-info → nextrec-0.2.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.2.4
+Version: 0.2.5
 Summary: A comprehensive recommendation library with match, ranking, and multi-task learning models
 Project-URL: Homepage, https://github.com/zerolovesea/NextRec
 Project-URL: Repository, https://github.com/zerolovesea/NextRec
@@ -61,7 +61,7 @@ Description-Content-Type: text/markdown
 ![Python](https://img.shields.io/badge/Python-3.10+-blue.svg)
 ![PyTorch](https://img.shields.io/badge/PyTorch-1.10+-ee4c2c.svg)
 ![License](https://img.shields.io/badge/License-Apache%202.0-green.svg)
-![Version](https://img.shields.io/badge/Version-0.2.4-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.2.5-orange.svg)
 
 English | [中文版](README_zh.md)
 

{nextrec-0.2.4.dist-info → nextrec-0.2.5.dist-info}/RECORD

@@ -1,21 +1,21 @@
 nextrec/__init__.py,sha256=CvocnY2uBp0cjNkhrT6ogw0q2bN9s1GNp754FLO-7lo,1117
-nextrec/__version__.py,sha256=SBl2EPFW-ltPvQ7vbVWItyAsz3aKYIpjO7vcfr84GkU,22
+nextrec/__version__.py,sha256=Xsa3ayOMVkhUWm4t06YeyHE0apjpZefxLH4ylp0CDtU,22
 nextrec/basic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nextrec/basic/activation.py,sha256=9EfYmwE0brTSKwx_0FIGQ_rybFBT9n_G-UWA7NAhMsI,2804
 nextrec/basic/callback.py,sha256=qkq3k8rP0g4BW2C3FSCdVt_CyCcJwJ-rUXjhT2p4LP8,1035
-nextrec/basic/features.py,sha256=TJfJgzWuy68lBKOeCzztcUK3ZtjHhK8oSMs8k0vXGlg,3961
-nextrec/basic/layers.py,sha256=mDNApSlPkmPSnIPj3BDHfDEjviLybWuSGrh61Zog2uk,38290
+nextrec/basic/features.py,sha256=pQyqs-hGFN8T6goTxIjaXxw9S5PG57OWX7EpsKFlb4c,4194
+nextrec/basic/layers.py,sha256=AezH_WYvU45eLw5EHmyAC69MCUkRTv_ZKbNW4WBC5iE,38071
 nextrec/basic/loggers.py,sha256=x8lzyyK-uqBN5XGOm1Cb33dmfc2bl114n6QeFTtE54k,3752
 nextrec/basic/metrics.py,sha256=w8tGe2tTbBNz9A1TNZF3jSpxcNC6QvFP5I0lWRd0Nw4,20398
-nextrec/basic/model.py,sha256=mtVttDWmrhdW-L1PAelJ90a1BW0q6bzG9roMvrPTU0U,66342
+nextrec/basic/model.py,sha256=k9dbV4CP-1wvr-QLJ0dF6nYtGugXQIFR6J8kZDS9iSs,63968
 nextrec/basic/session.py,sha256=2kogEjgKAN1_ygelbwoqOs187BAcUnDTqXG1w_Pgb9I,4791
 nextrec/data/__init__.py,sha256=HLnARJrqDEVPTcofPSAEimy2Oj15vbomj-7UvT4ze_4,767
-nextrec/data/data_utils.py,sha256=vGZ378YM_JQXO9npRB7JqojJx1ovjbJCWI-7lQJkicA,6298
-nextrec/data/dataloader.py,sha256=LAKpcSHhq53scq8PKwF8uqxa8wQLG0FshjY3TQwIvBU,20459
-nextrec/data/preprocessor.py,sha256=N7m4PYGZE6AND0XyYRvXKYAUub9aHGb1qmxbBRxlZKA,42294
-nextrec/loss/__init__.py,sha256=t-wkqxcu5wdYlrb67-CxX9aOGom0CpMJK8Fe8KGDSEE,857
+nextrec/data/data_utils.py,sha256=xz0xVBA7UzHXz7r_Yf0eMB5RrarPKg_1ZTdWvAqRZCM,7623
+nextrec/data/dataloader.py,sha256=vtgt2B7rUmIG7wg-HE2ZesBaD6cuS2PwklFCWGA9tCw,14142
+nextrec/data/preprocessor.py,sha256=J-3fo_LIz100spqCHoSpewYcneiZwhaCKyRdroPSjeY,41548
+nextrec/loss/__init__.py,sha256=mO5t417BneZ8Ysa51GyjDaffjWyjzFgPXIQrrggasaQ,827
 nextrec/loss/listwise.py,sha256=LcYIPf6PGRtjV_AoWaAyp3rse904S2MghE5t032I07I,5628
-nextrec/loss/loss_utils.py,sha256=LnTkpMTS2bhbq4Lsjf3AUn1uBaOg1TaH5VO2R8hwARc,5324
+nextrec/loss/loss_utils.py,sha256=cFNSvv-eaFwcfjLgxN3yNmf0L7ofC0ysgkUYjliLBpE,2535
 nextrec/loss/pairwise.py,sha256=RuQuTE-EkLaHQvT9m0CTAXxneTnVQLF1Pi9wblEClI8,3289
 nextrec/loss/pointwise.py,sha256=6QveizdohzQTxAoBKTVSoCBpp-fy3JC8vCjImXa7jL0,7157
 nextrec/models/generative/hstu.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -26,7 +26,7 @@ nextrec/models/match/dssm_v2.py,sha256=ywtqTy3YN9ke_7kzcDp7Fhtldw9RJz6yfewxALJb6
 nextrec/models/match/mind.py,sha256=XSUDlZ-V95JXHHBDUl5sz99SaVuQKDvf3TArVjwUexs,9417
 nextrec/models/match/sdm.py,sha256=96yfMQ6arP6JRhAkDTGEjlBiTteznMykrDV_3jqvvVk,10920
 nextrec/models/match/youtube_dnn.py,sha256=pnrz9LYu65Fj4neOriFF45B5k2-yYiiREtQICxxYXZ0,7546
-nextrec/models/multi_task/esmm.py,sha256=E9B6TlpnPUeyldTofyFg4B7SKByyxbiW2fUGHLOryO4,4883
+nextrec/models/multi_task/esmm.py,sha256=U9DkYxwAhD_uWB8H2tsx1zfIStp5Xp8bvM6sc8S9tu4,4889
 nextrec/models/multi_task/mmoe.py,sha256=zhQr43Vfz7Kgi6B9pKPmaenp_38a_D7w4VvlpwCyF6Y,6165
 nextrec/models/multi_task/ple.py,sha256=otP6oLgzrJhwkLFItzNE-AtIPouObDkafRvWzTCxfNo,11335
 nextrec/models/multi_task/share_bottom.py,sha256=LL5HBVlvvBzHV2fLBRQMGIwpqmlxILTgU4c51XyTCo4,4517
@@ -39,15 +39,16 @@ nextrec/models/ranking/dien.py,sha256=E6s9TDwQfGSwtzzh8hG2F5gwgVxzVZPcptYvHLNzOL
 nextrec/models/ranking/din.py,sha256=j5tkT5k91CbsMlMr5vJOySrcY2_rFGxmEgJJ0McW7-Q,7196
 nextrec/models/ranking/fibinet.py,sha256=X6CbQbritvq5jql_Tvs4bn_tRla2zpWPplftZv8k6f0,4853
 nextrec/models/ranking/fm.py,sha256=3Qx_Fgowegr6UPQtEeTmHtOrbWzkvqH94ZTjOqRLu-E,2961
-nextrec/models/ranking/masknet.py,sha256=hU3m270vd9DWH2_Hh1hYiCGaF9fKC3eIsWQLSA-Gdf8,12215
+nextrec/models/ranking/masknet.py,sha256=Tx5deIv7oShm4DdXX1IJL8Hz8-5uGqcPMK7pj00xTHg,12230
 nextrec/models/ranking/pnn.py,sha256=5RxIKdxD0XcGq-b_QDdwGRwk6b_5BQjyMvCw3Ibv2Kk,4957
 nextrec/models/ranking/widedeep.py,sha256=b6ctElaZPv5WSYDA4piYUBo3je0eJpWpWECwcuWavM4,3716
 nextrec/models/ranking/xdeepfm.py,sha256=I00J5tfE4tPluqeW-qrNtE4V_9fC7-rgFvA0Fxqka7o,4274
-nextrec/utils/__init__.py,sha256=6x3OZbqks2gtgJd00y_-Y8QiAT42x5t14ARHQ-ULQDo,350
+nextrec/utils/__init__.py,sha256=A3mH6M-DmDBWQ1stIIaTsNzvUy_AKaUWtRmrzU5R3FE,429
+nextrec/utils/common.py,sha256=-n4wSbP-EptpzLcJv6fV-ytBzPliOj6m-mrK_Qk6s4A,458
 nextrec/utils/embedding.py,sha256=yxYSdFx0cJITh3Gf-K4SdhwRtKGcI0jOsyBgZ0NLa_c,465
 nextrec/utils/initializer.py,sha256=ffYOs5QuIns_d_-5e40iNtg6s1ftgREJN-ueq_NbDQE,1647
 nextrec/utils/optimizer.py,sha256=85ifoy2IQgjPHOqLqr1ho7XBGE_0ry1yEB9efS6C2lM,2446
-nextrec-0.2.4.dist-info/METADATA,sha256=VurhzAYPQ_PbBi6WJFHvgbbk08OVd_1udwLPHxqApag,11425
-nextrec-0.2.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-nextrec-0.2.4.dist-info/licenses/LICENSE,sha256=2fQfVKeafywkni7MYHyClC6RGGC3laLTXCNBx-ubtp0,1064
-nextrec-0.2.4.dist-info/RECORD,,
+nextrec-0.2.5.dist-info/METADATA,sha256=Ya8KTj9x1ozIaciXYTKnTFBLGiC4buUBIz-jVHHAM3s,11425
+nextrec-0.2.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+nextrec-0.2.5.dist-info/licenses/LICENSE,sha256=2fQfVKeafywkni7MYHyClC6RGGC3laLTXCNBx-ubtp0,1064
+nextrec-0.2.5.dist-info/RECORD,,