torch-rechub 0.0.1__py3-none-any.whl → 0.0.4__py3-none-any.whl

torch_rechub/basic/loss_func.py

@@ -0,0 +1,223 @@
+import torch
+import torch.functional as F
+import torch.nn as nn
+
+
+class RegularizationLoss(nn.Module):
+    """Unified L1/L2 Regularization Loss for embedding and dense parameters.
+    
+    Example:
+        >>> reg_loss_fn = RegularizationLoss(embedding_l2=1e-5, dense_l2=1e-5)
+        >>> # In model's forward or trainer
+        >>> reg_loss = reg_loss_fn(model)
+        >>> total_loss = task_loss + reg_loss
+    """
+
+    def __init__(self, embedding_l1=0.0, embedding_l2=0.0, dense_l1=0.0, dense_l2=0.0):
+        super(RegularizationLoss, self).__init__()
+        self.embedding_l1 = embedding_l1
+        self.embedding_l2 = embedding_l2
+        self.dense_l1 = dense_l1
+        self.dense_l2 = dense_l2
+
+    def forward(self, model):
+        reg_loss = 0.0
+
+        # Register normalization layers
+        norm_params = set()
+        for module in model.modules():
+            if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.LayerNorm, nn.GroupNorm, nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d)):
+                for param in module.parameters():
+                    norm_params.add(id(param))
+
+        # Register embedding layers
+        embedding_params = set()
+        for module in model.modules():
+            if isinstance(module, (nn.Embedding, nn.EmbeddingBag)):
+                for param in module.parameters():
+                    embedding_params.add(id(param))
+
+        for param in model.parameters():
+            if param.requires_grad:
+                # Skip normalization layer parameters
+                if id(param) in norm_params:
+                    continue
+
+                if id(param) in embedding_params:
+                    if self.embedding_l1 > 0:
+                        reg_loss += self.embedding_l1 * torch.sum(torch.abs(param))
+                    if self.embedding_l2 > 0:
+                        reg_loss += self.embedding_l2 * torch.sum(param**2)
+                else:
+                    if self.dense_l1 > 0:
+                        reg_loss += self.dense_l1 * torch.sum(torch.abs(param))
+                    if self.dense_l2 > 0:
+                        reg_loss += self.dense_l2 * torch.sum(param**2)
+
+        return reg_loss
+
+
+class HingeLoss(torch.nn.Module):
+    """Hinge Loss for pairwise learning.
+    reference: https://github.com/ustcml/RecStudio/blob/main/recstudio/model/loss_func.py
+
+    """
+
+    def __init__(self, margin=2, num_items=None):
+        super().__init__()
+        self.margin = margin
+        self.n_items = num_items
+
+    def forward(self, pos_score, neg_score):
+        loss = torch.maximum(torch.max(neg_score, dim=-1).values - pos_score + self.margin, torch.tensor([0]).type_as(pos_score))
+        if self.n_items is not None:
+            impostors = neg_score - pos_score.view(-1, 1) + self.margin > 0
+            rank = torch.mean(impostors, -1) * self.n_items
+            return torch.mean(loss * torch.log(rank + 1))
+        else:
+            return torch.mean(loss)
+
+
+class BPRLoss(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, pos_score, neg_score):
+        loss = torch.mean(-(pos_score - neg_score).sigmoid().log(), dim=-1)
+        return loss
+
+
+class NCELoss(torch.nn.Module):
+    """Noise Contrastive Estimation (NCE) Loss for recommendation systems.
+
+    NCE Loss is more efficient than CrossEntropyLoss for large-scale recommendation
+    scenarios. It uses in-batch negatives to reduce computational complexity.
+
+    Reference:
+        - Noise-contrastive estimation: A new estimation principle for unnormalized
+          statistical models (Gutmann & Hyvärinen, 2010)
+        - HLLM: Hierarchical Large Language Model for Recommendation
+
+    Args:
+        temperature (float): Temperature parameter for scaling logits. Default: 1.0
+        ignore_index (int): Index to ignore in loss computation. Default: 0
+        reduction (str): Specifies the reduction to apply to the output.
+                        Options: 'mean', 'sum', 'none'. Default: 'mean'
+
+    Example:
+        >>> nce_loss = NCELoss(temperature=0.1)
+        >>> logits = torch.randn(32, 1000)  # (batch_size, vocab_size)
+        >>> targets = torch.randint(0, 1000, (32,))
+        >>> loss = nce_loss(logits, targets)
+    """
+
+    def __init__(self, temperature=1.0, ignore_index=0, reduction='mean'):
+        super().__init__()
+        self.temperature = temperature
+        self.ignore_index = ignore_index
+        self.reduction = reduction
+
+    def forward(self, logits, targets):
+        """Compute NCE loss.
+
+        Args:
+            logits (torch.Tensor): Model output logits of shape (batch_size, vocab_size)
+            targets (torch.Tensor): Target indices of shape (batch_size,)
+
+        Returns:
+            torch.Tensor: NCE loss value
+        """
+        # Scale logits by temperature
+        logits = logits / self.temperature
+
+        # Compute log softmax
+        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+
+        # Get log probability of target class
+        batch_size = targets.shape[0]
+        target_log_probs = log_probs[torch.arange(batch_size), targets]
+
+        # Create mask for ignore_index
+        mask = targets != self.ignore_index
+
+        # Compute loss
+        loss = -target_log_probs
+
+        # Apply mask
+        if mask.any():
+            loss = loss[mask]
+
+        # Apply reduction
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        else:  # 'none'
+            return loss
+
+
+class InBatchNCELoss(torch.nn.Module):
+    """In-Batch NCE Loss with explicit negative sampling.
+
+    This loss function uses other samples in the batch as negative samples,
+    which is more efficient than sampling random negatives.
+
+    Args:
+        temperature (float): Temperature parameter for scaling logits. Default: 0.1
+        ignore_index (int): Index to ignore in loss computation. Default: 0
+        reduction (str): Specifies the reduction to apply to the output.
+                        Options: 'mean', 'sum', 'none'. Default: 'mean'
+
+    Example:
+        >>> loss_fn = InBatchNCELoss(temperature=0.1)
+        >>> embeddings = torch.randn(32, 256)  # (batch_size, embedding_dim)
+        >>> item_embeddings = torch.randn(1000, 256)  # (vocab_size, embedding_dim)
+        >>> targets = torch.randint(0, 1000, (32,))
+        >>> loss = loss_fn(embeddings, item_embeddings, targets)
+    """
+
+    def __init__(self, temperature=0.1, ignore_index=0, reduction='mean'):
+        super().__init__()
+        self.temperature = temperature
+        self.ignore_index = ignore_index
+        self.reduction = reduction
+
+    def forward(self, embeddings, item_embeddings, targets):
+        """Compute in-batch NCE loss.
+
+        Args:
+            embeddings (torch.Tensor): User/query embeddings of shape (batch_size, embedding_dim)
+            item_embeddings (torch.Tensor): Item embeddings of shape (vocab_size, embedding_dim)
+            targets (torch.Tensor): Target item indices of shape (batch_size,)
+
+        Returns:
+            torch.Tensor: In-batch NCE loss value
+        """
+        # Compute logits: (batch_size, vocab_size)
+        logits = torch.matmul(embeddings, item_embeddings.t()) / self.temperature
+
+        # Compute log softmax
+        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+
+        # Get log probability of target class
+        batch_size = targets.shape[0]
+        target_log_probs = log_probs[torch.arange(batch_size), targets]
+
+        # Create mask for ignore_index
+        mask = targets != self.ignore_index
+
+        # Compute loss
+        loss = -target_log_probs
+
+        # Apply mask
+        if mask.any():
+            loss = loss[mask]
+
+        # Apply reduction
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        else:  # 'none'
+            return loss

torch_rechub/basic/metaoptimizer.py

@@ -0,0 +1,76 @@
+"""The metaoptimizer module, it provides a class MetaBalance
+MetaBalance is used to scale the gradient and balance the gradient of each task
+Authors: Qida Dong, dongjidan@126.com
+"""
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class MetaBalance(Optimizer):
+    """MetaBalance Optimizer
+       This method is used to scale the gradient and balance the gradient of each task
+
+    Args:
+        parameters (list): the parameters of model
+        relax_factor (float, optional): the relax factor of gradient scaling (default: 0.7)
+        beta (float, optional): the coefficient of moving average (default: 0.9)
+                """
+
+    def __init__(self, parameters, relax_factor=0.7, beta=0.9):
+
+        if relax_factor < 0. or relax_factor >= 1.:
+            raise ValueError(f'Invalid relax_factor: {relax_factor}, it should be 0. <= relax_factor < 1.')
+        if beta < 0. or beta >= 1.:
+            raise ValueError(f'Invalid beta: {beta}, it should be 0. <= beta < 1.')
+        rel_beta_dict = {'relax_factor': relax_factor, 'beta': beta}
+        super(MetaBalance, self).__init__(parameters, rel_beta_dict)
+
+    @torch.no_grad()
+    def step(self, losses):
+        """_summary_
+        Args:
+            losses (_type_): _description_
+
+        Raises:
+            RuntimeError: _description_
+        """
+
+        for idx, loss in enumerate(losses):
+            loss.backward(retain_graph=True)
+            for group in self.param_groups:
+                for gp in group['params']:
+                    if gp.grad is None:
+                        # print('breaking')
+                        break
+                    if gp.grad.is_sparse:
+                        raise RuntimeError('MetaBalance does not support sparse gradients')
+# store the result of moving average
+                    state = self.state[gp]
+                    if len(state) == 0:
+                        for i in range(len(losses)):
+                            if i == 0:
+                                gp.norms = [0]
+                            else:
+                                gp.norms.append(0)
+
+
+# calculate the moving average
+                    beta = group['beta']
+                    gp.norms[idx] = gp.norms[idx] * beta + \
+                        (1 - beta) * torch.norm(gp.grad)
+                    # scale the auxiliary gradient
+                    relax_factor = group['relax_factor']
+                    gp.grad = gp.grad * \
+                        gp.norms[0] / (gp.norms[idx] + 1e-5) * relax_factor + gp.grad * (1. - relax_factor)
+                    # store the gradient of each auxiliary task in state
+                    if idx == 0:
+                        state['sum_gradient'] = torch.zeros_like(gp.data)
+                        state['sum_gradient'] += gp.grad
+                    else:
+                        state['sum_gradient'] += gp.grad
+
+                    if gp.grad is not None:
+                        gp.grad.detach_()
+                        gp.grad.zero_()
+                    if idx == len(losses) - 1:
+                        gp.grad = state['sum_gradient']

torch_rechub/basic/metric.py

@@ -0,0 +1,251 @@
+"""The metric module, it is used to provide some metrics for recommenders.
+Available function:
+- auc_score: compute AUC
+- gauc_score: compute GAUC
+- log_loss: compute LogLoss
+- topk_metrics: compute topk metrics contains 'ndcg', 'mrr', 'recall', 'hit'
+Authors: Qida Dong, dongjidan@126.com
+"""
+from collections import defaultdict
+
+import numpy as np
+from sklearn.metrics import roc_auc_score
+
+
+def auc_score(y_true, y_pred):
+
+    return roc_auc_score(y_true, y_pred)
+
+
+def get_user_pred(y_true, y_pred, users):
+    """divide the result into different group by user id
+
+        Args:
+                y_true (array): all true labels of the data
+                y_pred (array): the predicted score
+                users (array): user id
+
+        Return:
+                user_pred (dict): {userid: values}, key is user id and value is the labels and scores of each user
+        """
+    user_pred = {}
+    for i, u in enumerate(users):
+        if u not in user_pred:
+            user_pred[u] = {'y_true': [y_true[i]], 'y_pred': [y_pred[i]]}
+        else:
+            user_pred[u]['y_true'].append(y_true[i])
+            user_pred[u]['y_pred'].append(y_pred[i])
+
+    return user_pred
+
+
+def gauc_score(y_true, y_pred, users, weights=None):
+    """compute GAUC
+
+        Args:
+                y_true (array): dim(N, ), all true labels of the data
+                y_pred (array): dim(N, ), the predicted score
+                users (array): dim(N, ), user id
+                weight (dict): {userid: weight_value}, it contains weights for each group.
+                                if it is None, the weight is equal to the number
+                                of times the user is recommended
+        Return:
+                score: float, GAUC
+        """
+    assert len(y_true) == len(y_pred) and len(y_true) == len(users)
+
+    user_pred = get_user_pred(y_true, y_pred, users)
+    score = 0
+    num = 0
+    for u in user_pred.keys():
+        auc = auc_score(user_pred[u]['y_true'], user_pred[u]['y_pred'])
+        if weights is None:
+            user_weight = len(user_pred[u]['y_true'])
+        else:
+            user_weight = weights[u]
+        auc *= user_weight
+        num += user_weight
+        score += auc
+    return score / num
+
+
+def ndcg_score(y_true, y_pred, topKs=None):
+    if topKs is None:
+        topKs = [5]
+    result = topk_metrics(y_true, y_pred, topKs)
+    return result['NDCG']
+
+
+def hit_score(y_true, y_pred, topKs=None):
+    if topKs is None:
+        topKs = [5]
+    result = topk_metrics(y_true, y_pred, topKs)
+    return result['Hit']
+
+
+def mrr_score(y_true, y_pred, topKs=None):
+    if topKs is None:
+        topKs = [5]
+    result = topk_metrics(y_true, y_pred, topKs)
+    return result['MRR']
+
+
+def recall_score(y_true, y_pred, topKs=None):
+    if topKs is None:
+        topKs = [5]
+    result = topk_metrics(y_true, y_pred, topKs)
+    return result['Recall']
+
+
+def precision_score(y_true, y_pred, topKs=None):
+    if topKs is None:
+        topKs = [5]
+    result = topk_metrics(y_true, y_pred, topKs)
+    return result['Precision']
+
+
+def topk_metrics(y_true, y_pred, topKs=None):
+    """choice topk metrics and compute it
+        the metrics contains 'ndcg', 'mrr', 'recall', 'precision' and 'hit'
+
+        Args:
+                y_true (dict): {userid, item_ids}, the key is user id and the value is the list that contains the items the user interacted
+                y_pred (dict): {userid, item_ids}, the key is user id and the value is the list that contains the items recommended
+                topKs (list or tuple): if you want to get top5 and top10, topKs=(5, 10)
+
+        Return:
+                results (dict): {metric_name: metric_values}, it contains five metrics, 'ndcg', 'recall', 'mrr', 'hit', 'precision'
+
+        """
+    if topKs is None:
+        topKs = [5]
+    assert len(y_true) == len(y_pred)
+
+    if not isinstance(topKs, (tuple, list)):
+        raise ValueError('topKs wrong, it should be tuple or list')
+
+    pred_array = []
+    true_array = []
+    for u in y_true.keys():
+        pred_array.append(y_pred[u])
+        true_array.append(y_true[u])
+    ndcg_result = []
+    mrr_result = []
+    hit_result = []
+    precision_result = []
+    recall_result = []
+    for idx in range(len(topKs)):
+        ndcgs = 0
+        mrrs = 0
+        hits = 0
+        precisions = 0
+        recalls = 0
+        gts = 0
+        for i in range(len(true_array)):
+            if len(true_array[i]) != 0:
+                mrr_tmp = 0
+                mrr_flag = True
+                hit_tmp = 0
+                dcg_tmp = 0
+                idcg_tmp = 0
+                for j in range(topKs[idx]):
+                    if pred_array[i][j] in true_array[i]:
+                        hit_tmp += 1.
+                        if mrr_flag:
+                            mrr_flag = False
+                            mrr_tmp = 1. / (1 + j)
+                        dcg_tmp += 1. / (np.log2(j + 2))
+                    if j < len(true_array[i]):
+                        idcg_tmp += 1. / (np.log2(j + 2))
+                gts += len(true_array[i])
+                hits += hit_tmp
+                mrrs += mrr_tmp
+                recalls += hit_tmp / len(true_array[i])
+                precisions += hit_tmp / topKs[idx]
+                if idcg_tmp != 0:
+                    ndcgs += dcg_tmp / idcg_tmp
+        hit_result.append(round(hits / gts, 4))
+        mrr_result.append(round(mrrs / len(pred_array), 4))
+        recall_result.append(round(recalls / len(pred_array), 4))
+        precision_result.append(round(precisions / len(pred_array), 4))
+        ndcg_result.append(round(ndcgs / len(pred_array), 4))
+
+    results = defaultdict(list)
+    for idx in range(len(topKs)):
+
+        output = f'NDCG@{topKs[idx]}: {ndcg_result[idx]}'
+        results['NDCG'].append(output)
+
+        output = f'MRR@{topKs[idx]}: {mrr_result[idx]}'
+        results['MRR'].append(output)
+
+        output = f'Recall@{topKs[idx]}: {recall_result[idx]}'
+        results['Recall'].append(output)
+
+        output = f'Hit@{topKs[idx]}: {hit_result[idx]}'
+        results['Hit'].append(output)
+
+        output = f'Precision@{topKs[idx]}: {precision_result[idx]}'
+        results['Precision'].append(output)
+    return results
+
+
+def log_loss(y_true, y_pred):
+    score = y_true * np.log(y_pred) + (1 - y_true) * np.log(1 - y_pred)
+    return -score.sum() / len(y_true)
+
+
+def Coverage(y_pred, all_items, topKs=None):
+    """compute the coverage
+        This method measures the diversity of the recommended items
+        and the ability to explore the long-tailed items
+        Arg:
+                y_pred (dict): {userid, item_ids}, the key is user id and the value is the list that contains the items recommended
+                all_items (set): all unique items
+        Return:
+                result (list[float]): the list of coverage scores
+        """
+    if topKs is None:
+        topKs = [5]
+    result = []
+    for k in topKs:
+        rec_items = set([])
+        for u in y_pred.keys():
+            tmp_items = set(y_pred[u][:k])
+            rec_items = rec_items | tmp_items
+        score = len(rec_items) * 1. / len(all_items)
+        score = round(score, 4)
+        result.append(f'Coverage@{k}: {score}')
+    return result
+
+
+# print(Coverage({'0':[0,1,2],'1':[1,3,4]}, {0,1,2,3,4,5}, [2,3]))
+
+# pred = np.array([  0.3, 0.2, 0.5, 0.9, 0.7, 0.31, 0.8, 0.1, 0.4, 0.6])
+# label = np.array([   1,   0,   0,   1,   0,   0,    1,   0,   0,   1])
+# users_id = np.array([ 2,   1,   0,   2,   1,   0,    0,   2,   1,   1])
+
+# print('auc: ', auc_score(label, pred))
+# print('gauc: ', gauc_score(label, pred, users_id))
+# print('log_loss: ', log_loss(label, pred))
+
+# for mt in ['ndcg', 'mrr', 'recall', 'hit','s']:
+# 	tm = topk_metrics(y_true, y_pred, users_id, 3, metric_type=mt)
+# 	print(f'{mt}: {tm}')
+# y_pred = {'0': [0, 1], '1': [0, 1], '2': [2, 3]}
+# y_true = {'0': [1, 2], '1': [0, 1, 2], '2': [2, 3]}
+# out = topk_metrics(y_true, y_pred, topKs=(1,2))
+# ndcgs = ndcg_score(y_true,y_pred, topKs=(1,2))
+# print(out)
+# print(ndcgs)
+
+# ground_truth, match_res = np.load("C:\\Users\\dongj\\Desktop/res.npy", allow_pickle=True)
+# print(len(ground_truth),len(match_res))
+# out = topk_metrics(y_true=ground_truth, y_pred=match_res, topKs=[50])
+# print(out)
+
+if __name__ == "__main__":
+    y_pred = {'0': [0, 1], '1': [0, 1], '2': [2, 3]}
+    y_true = {'0': [1, 2], '1': [0, 1, 2], '2': [2, 3]}
+    out = topk_metrics(y_true, y_pred, topKs=(1, 2))
+    print(out)

torch_rechub/models/generative/__init__.py

@@ -0,0 +1,6 @@
+"""Generative Recommendation Models."""
+
+from .hllm import HLLMModel
+from .hstu import HSTUModel
+
+__all__ = ['HSTUModel', 'HLLMModel']