nextrec 0.3.6__py3-none-any.whl → 0.4.1__py3-none-any.whl

nextrec/models/ranking/xdeepfm.py

@@ -56,7 +56,7 @@ class xDeepFM(BaseModel):
         return "xDeepFM"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
     
     def __init__(self,
@@ -67,6 +67,7 @@ class xDeepFM(BaseModel):
                  cin_size: list[int] = [128, 128],
                  split_half: bool = True,
                  target: list[str] = [],
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | None = "bce",
@@ -83,13 +84,12 @@ class xDeepFM(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -118,7 +118,7 @@ class xDeepFM(BaseModel):
         deep_emb_dim_total = sum([f.embedding_dim for f in self.deep_features if not isinstance(f, DenseFeature)])
         dense_input_dim = sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
         self.mlp = MLP(input_dim=deep_emb_dim_total + dense_input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.task)
 
         # Register regularization weights
         self.register_regularization_weights(

nextrec/utils/__init__.py

@@ -12,10 +12,10 @@ This package provides various utility functions organized by category:
 - feature_utils: Feature processing utilities
 
 Date: create on 13/11/2025
-Last update: 03/12/2025 (refactored)
+Last update: 06/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
-
+from . import optimizer, initializer, embedding
 from .optimizer import get_optimizer, get_scheduler
 from .initializer import get_initializer
 from .embedding import get_auto_embedding_dim
@@ -24,7 +24,7 @@ from .tensor import to_tensor, stack_tensors, concat_tensors, pad_sequence_tenso
 from .file import resolve_file_paths, read_table, load_dataframes, iter_file_chunks, default_output_dir
 from .model import merge_features, get_mlp_output_dim
 from .feature import normalize_to_list
-from . import optimizer, initializer, embedding
+from .synthetic_data import generate_ranking_data, generate_distributed_ranking_data, generate_match_data, generate_multitask_data
 
 __all__ = [
     # Optimizer & Scheduler
@@ -61,6 +61,10 @@ __all__ = [
     # Feature utilities
     'normalize_to_list',
     
+    # Synthetic data utilities
+    'generate_ranking_data',
+    'generate_distributed_ranking_data',
+    
     # Module exports
     'optimizer',
     'initializer',

nextrec/utils/device.py

@@ -7,6 +7,7 @@ Author: Yang Zhou, zyaztec@gmail.com
 import os
 import torch
 import platform
+import logging
 import multiprocessing
 
 
@@ -36,3 +37,32 @@ def get_device_info() -> dict:
         info['cuda_capability'] = torch.cuda.get_device_capability(0)
     
     return info
+
+def configure_device(
+    distributed: bool,
+    local_rank: int,
+    base_device: torch.device | str = "cpu"
+) -> torch.device:
+    try:
+        device = torch.device(base_device)
+    except Exception:
+        logging.warning("[configure_device Warning] Invalid base_device, falling back to CPU.")
+        return torch.device("cpu")
+
+    if distributed:
+        if device.type == "cuda":
+            if not torch.cuda.is_available():
+                logging.warning("[Distributed Warning] CUDA requested but unavailable. Falling back to CPU.")
+                return torch.device("cpu")
+            if not (0 <= local_rank < torch.cuda.device_count()):
+                logging.warning(f"[Distributed Warning] local_rank {local_rank} is invalid for available CUDA devices. Falling back to CPU.")
+                return torch.device("cpu")
+            try:
+                torch.cuda.set_device(local_rank)
+                return torch.device(f"cuda:{local_rank}")
+            except Exception as exc:
+                logging.warning(f"[Distributed Warning] Failed to set CUDA device for local_rank {local_rank}: {exc}. Falling back to CPU.")
+                return torch.device("cpu")
+        else:
+            return torch.device("cpu")
+    return device

nextrec/utils/distributed.py

@@ -0,0 +1,114 @@
+"""
+Distributed utilities for NextRec.
+
+Date: create on 04/12/2025
+Checkpoint: edit on 05/12/2025
+Author: Yang Zhou,zyaztec@gmail.com
+"""
+
+import logging
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from torch.utils.data import DataLoader, IterableDataset
+from torch.utils.data.distributed import DistributedSampler
+from nextrec.basic.loggers import colorize
+
+def init_process_group(distributed: bool, rank: int, world_size: int, device_id: int | None = None) -> None:
+    """
+    initialize distributed process group for multi-GPU training.
+    
+    Args:
+        distributed: whether to enable distributed training
+        rank: global rank of the current process
+        world_size: total number of processes
+    """
+    if (not distributed) or (not dist.is_available()) or dist.is_initialized():
+        return
+    backend = "nccl" if device_id is not None else "gloo"
+    if backend == "nccl":
+        torch.cuda.set_device(device_id)
+    dist.init_process_group(backend=backend, init_method="env://", rank=rank, world_size=world_size)
+
+def gather_numpy(self, array: np.ndarray | None) -> np.ndarray | None:
+    """
+    Gather numpy arrays (or None) across ranks. Uses all_gather_object to avoid
+    shape mismatches and ensures every rank participates even when local data is empty.
+    """
+    if not (self.distributed and dist.is_available() and dist.is_initialized()):
+        return array
+
+    world_size = dist.get_world_size()
+    gathered: list[np.ndarray | None] = [None for _ in range(world_size)]
+    dist.all_gather_object(gathered, array)
+    pieces: list[np.ndarray] = []
+    for item in gathered:
+        if item is None:
+            continue
+        item_np = np.asarray(item)
+        if item_np.size > 0:
+            pieces.append(item_np)
+    if not pieces:
+        return None
+    return np.concatenate(pieces, axis=0)
+
+def add_distributed_sampler(
+    loader: DataLoader,
+    distributed: bool,
+    world_size: int,
+    rank: int,
+    shuffle: bool,
+    drop_last: bool,
+    default_batch_size: int,
+    is_main_process: bool = False,
+) -> tuple[DataLoader, DistributedSampler | None]:
+    """
+    add distributedsampler to a dataloader, this for distributed training 
+    when each device has its own dataloader
+    """
+    # early return if not distributed
+    if not (distributed and dist.is_available() and dist.is_initialized()):
+        return loader, None
+    # return if already has DistributedSampler
+    if isinstance(loader.sampler, DistributedSampler):
+        return loader, loader.sampler
+    dataset = getattr(loader, "dataset", None)
+    if dataset is None:
+        return loader, None
+    if isinstance(dataset, IterableDataset):
+        if is_main_process:
+            logging.info(colorize("[Distributed Info] Iterable/streaming DataLoader provided; DistributedSampler is skipped. Ensure dataset handles sharding per rank.", color="yellow"))
+        return loader, None
+    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=shuffle, drop_last=drop_last)
+    loader_kwargs = {
+        "batch_size": loader.batch_size if loader.batch_size is not None else default_batch_size,
+        "shuffle": False,
+        "sampler": sampler,
+        "num_workers": loader.num_workers,
+        "collate_fn": loader.collate_fn,
+        "drop_last": drop_last,
+    }
+    if getattr(loader, "pin_memory", False):
+        loader_kwargs["pin_memory"] = True
+    pin_memory_device = getattr(loader, "pin_memory_device", None)
+    if pin_memory_device:
+        loader_kwargs["pin_memory_device"] = pin_memory_device
+    timeout = getattr(loader, "timeout", None)
+    if timeout:
+        loader_kwargs["timeout"] = timeout
+    worker_init_fn = getattr(loader, "worker_init_fn", None)
+    if worker_init_fn is not None:
+        loader_kwargs["worker_init_fn"] = worker_init_fn
+    generator = getattr(loader, "generator", None)
+    if generator is not None:
+        loader_kwargs["generator"] = generator
+    if loader.num_workers > 0:
+        loader_kwargs["persistent_workers"] = getattr(loader, "persistent_workers", False)
+        prefetch_factor = getattr(loader, "prefetch_factor", None)
+        if prefetch_factor is not None:
+            loader_kwargs["prefetch_factor"] = prefetch_factor
+    distributed_loader = DataLoader(dataset, **loader_kwargs)
+    if is_main_process:
+        logging.info(colorize("[Distributed Info] Attached DistributedSampler to provided DataLoader", color="cyan"))
+    return distributed_loader, sampler

nextrec/utils/synthetic_data.py

@@ -0,0 +1,413 @@
+"""
+Synthetic Data Generation Utilities
+
+This module provides utilities for generating synthetic datasets for testing
+and tutorial purposes in the NextRec framework.
+
+Date: create on 06/12/2025
+Author: Yang Zhou, zyaztec@gmail.com
+"""
+
+import numpy as np
+import pandas as pd
+from typing import Optional, Dict, List, Tuple, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
+
+def generate_ranking_data(
+    n_samples: int = 10000,
+    n_dense: int = 5,
+    n_sparse: int = 8,
+    n_sequences: int = 2,
+    user_vocab_size: int = 1000,
+    item_vocab_size: int = 500,
+    sparse_vocab_size: int = 50,
+    sequence_max_len: int = 20,
+    embedding_dim: int = 16,
+    seed: int = 42,
+    custom_sparse_features: Optional[Dict[str, int]] = None,
+    use_simple_names: bool = True
+) -> Tuple[pd.DataFrame, List, List, List]:
+    """
+    Generate synthetic data for ranking tasks (CTR prediction)
+    
+    Returns:
+        tuple: (dataframe, dense_features, sparse_features, sequence_features)
+    """
+    print(f"Generating {n_samples} synthetic ranking samples...")
+    
+    np.random.seed(seed)
+    data = {}
+    
+    for i in range(n_dense):
+        data[f'dense_{i}'] = np.random.randn(n_samples).astype(np.float32)
+    
+    # Generate basic sparse features (always include user_id and item_id)
+    data['user_id'] = np.random.randint(1, user_vocab_size, n_samples)
+    data['item_id'] = np.random.randint(1, item_vocab_size, n_samples)
+    
+    # Generate additional sparse features
+    if custom_sparse_features:
+        for feat_name, vocab_size in custom_sparse_features.items():
+            data[feat_name] = np.random.randint(0, vocab_size, n_samples)
+    else:
+        for i in range(n_sparse - 2):
+            data[f'sparse_{i}'] = np.random.randint(1, sparse_vocab_size, n_samples)
+    
+    # Generate sequence features (list of IDs)
+    sequence_names = []
+    sequence_vocabs = []
+    
+    for i in range(n_sequences):
+        sequences = []
+        for _ in range(n_samples):
+            seq_len = np.random.randint(5, sequence_max_len + 1)
+            if i == 0:
+                # First sequence uses item vocabulary
+                seq = np.random.randint(0, item_vocab_size, seq_len).tolist()
+                seq_vocab = item_vocab_size
+                if custom_sparse_features:
+                    seq_name = 'hist_items'
+                else:
+                    seq_name = 'sequence_0'
+            else:
+                # Other sequences use category vocabulary
+                if custom_sparse_features and 'category' in custom_sparse_features:
+                    seq_vocab = custom_sparse_features['category']
+                    seq = np.random.randint(0, seq_vocab, seq_len).tolist()
+                    seq_name = f'hist_categories' if i == 1 else f'sequence_{i}'
+                else:
+                    seq_vocab = sparse_vocab_size
+                    seq = np.random.randint(0, seq_vocab, seq_len).tolist()
+                    seq_name = f'sequence_{i}'
+            
+            # Padding
+            seq = seq + [0] * (sequence_max_len - len(seq))
+            sequences.append(seq)
+        
+        data[seq_name] = sequences
+        sequence_names.append(seq_name)
+        sequence_vocabs.append(seq_vocab)
+    
+    if 'gender' in data and 'dense_0' in data:
+        # Complex label generation with feature correlation
+        label_probs = 1 / (1 + np.exp(-(
+            data['dense_0'] * 0.3 +
+            data['dense_1'] * 0.2 +
+            (data['gender'] - 0.5) * 0.5 +
+            np.random.randn(n_samples) * 0.1
+        )))
+        data['label'] = (label_probs > 0.5).astype(np.float32)
+    else:
+        data['label'] = np.random.randint(0, 2, n_samples).astype(np.float32)
+    
+    df = pd.DataFrame(data)
+    print(f"Generated data shape: {df.shape}")
+    if 'gender' in data:
+        print(f"Positive rate: {data['label'].mean():.4f}")
+    
+    # Import here to avoid circular import
+    from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
+    
+    # Create feature definitions
+    # Use input_dim for dense features to be compatible with both simple and complex scenarios
+    dense_features = [DenseFeature(name=f'dense_{i}', input_dim=1)  for i in range(n_dense)]
+    
+    # Create sparse features
+    sparse_features = [SparseFeature( name='user_id',  embedding_name='user_emb',  vocab_size=user_vocab_size, embedding_dim=embedding_dim),
+                       SparseFeature(name='item_id', embedding_name='item_emb', vocab_size=item_vocab_size,  embedding_dim=embedding_dim),]
+    
+    if custom_sparse_features:
+        # Add custom sparse features with proper vocab sizes
+        for feat_name, vocab_size in custom_sparse_features.items():
+            sparse_features.append(SparseFeature(name=feat_name, embedding_name=f'{feat_name}_emb', vocab_size=vocab_size, embedding_dim=embedding_dim))
+    else:
+        # Add generic sparse features
+        sparse_features.extend([SparseFeature(name=f'sparse_{i}',  embedding_name=f'sparse_{i}_emb', vocab_size=sparse_vocab_size,  embedding_dim=embedding_dim) for i in range(n_sparse - 2)])
+    
+    # Create sequence features
+    sequence_features = []
+    for i, (seq_name, seq_vocab) in enumerate(zip(sequence_names, sequence_vocabs)):
+        if i == 0:
+            # First sequence shares embedding with item_id
+            embedding_name = 'item_emb'
+        elif custom_sparse_features and 'category' in custom_sparse_features and seq_name == 'hist_categories':
+            # hist_categories shares embedding with category
+            embedding_name = 'category_emb'
+        else:
+            # Other sequences share with sparse_0
+            embedding_name = 'sparse_0_emb'
+        sequence_features.append(SequenceFeature(name=seq_name, vocab_size=seq_vocab, max_len=sequence_max_len, embedding_dim=embedding_dim, padding_idx=0, embedding_name=embedding_name))
+    return df, dense_features, sparse_features, sequence_features
+
+
+def generate_match_data(
+    n_samples: int = 10000,
+    user_vocab_size: int = 1000,
+    item_vocab_size: int = 5000,
+    category_vocab_size: int = 100,
+    brand_vocab_size: int = 200,
+    city_vocab_size: int = 100,
+    user_feature_vocab_size: int = 50,
+    item_feature_vocab_size: int = 50,
+    sequence_max_len: int = 50,
+    user_embedding_dim: int = 32,
+    item_embedding_dim: int = 32,
+    seed: int = 42
+) -> Tuple[pd.DataFrame, List, List, List, List, List, List]:
+    """
+    Generate synthetic data for match/retrieval tasks
+    
+    Returns:
+        tuple: (dataframe, user_dense_features, user_sparse_features, user_sequence_features,
+                item_dense_features, item_sparse_features, item_sequence_features)
+    """
+    print(f"Generating {n_samples} synthetic match samples...")
+    
+    np.random.seed(seed)
+    data = {}
+    
+    # User features
+    data['user_id'] = np.random.randint(1, user_vocab_size, n_samples)
+    data['user_age'] = np.random.randn(n_samples).astype(np.float32)
+    data['user_gender'] = np.random.randint(0, 2, n_samples)
+    data['user_city'] = np.random.randint(0, city_vocab_size, n_samples)
+    
+    for i in range(3):
+        data[f'user_feature_{i}'] = np.random.randint(1, user_feature_vocab_size, n_samples)
+    
+    # User behavior sequences
+    user_hist_items = []
+    user_hist_categories = []
+    for _ in range(n_samples):
+        seq_len = np.random.randint(10, sequence_max_len + 1)
+        hist_items = np.random.randint(1, item_vocab_size, seq_len).tolist()
+        hist_items = hist_items + [0] * (sequence_max_len - len(hist_items))
+        user_hist_items.append(hist_items)
+        
+        hist_cats = np.random.randint(1, category_vocab_size, seq_len).tolist()
+        hist_cats = hist_cats + [0] * (sequence_max_len - len(hist_cats))
+        user_hist_categories.append(hist_cats)
+    
+    data['user_hist_items'] = user_hist_items
+    data['user_hist_categories'] = user_hist_categories
+    
+    # Item features
+    data['item_id'] = np.random.randint(1, item_vocab_size, n_samples)
+    data['item_price'] = np.random.randn(n_samples).astype(np.float32)
+    data['item_category'] = np.random.randint(1, category_vocab_size, n_samples)
+    data['item_brand'] = np.random.randint(1, brand_vocab_size, n_samples)
+    
+    for i in range(3):
+        data[f'item_feature_{i}'] = np.random.randint(1, item_feature_vocab_size, n_samples)
+    
+    # Generate labels with some correlation to features
+    label_probs = 1 / (1 + np.exp(-(
+        data['user_age'] * 0.2 +
+        (data['user_gender'] - 0.5) * 0.3 +
+        data['item_price'] * 0.15 +
+        np.random.randn(n_samples) * 0.5
+    )))
+    data['label'] = (label_probs > 0.5).astype(np.float32)
+    
+    df = pd.DataFrame(data)
+    print(f"Generated data shape: {df.shape}")
+    print(f"Positive rate: {data['label'].mean():.4f}")
+    
+    # Import here to avoid circular import
+    from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
+    
+    # User dense features
+    user_dense_features = [DenseFeature(name='user_age', input_dim=1)]
+    
+    # User sparse features
+    user_sparse_features = [
+        SparseFeature(name='user_id', vocab_size=user_vocab_size, embedding_dim=user_embedding_dim),
+        SparseFeature(name='user_gender', vocab_size=2, embedding_dim=8),
+        SparseFeature(name='user_city', vocab_size=city_vocab_size, embedding_dim=16),
+    ]
+    user_sparse_features.extend([
+        SparseFeature(name=f'user_feature_{i}', vocab_size=user_feature_vocab_size, embedding_dim=8)
+        for i in range(3)
+    ])
+    
+    # User sequence features
+    user_sequence_features = [
+        SequenceFeature(name='user_hist_items', vocab_size=item_vocab_size, 
+                       max_len=sequence_max_len, embedding_dim=user_embedding_dim, padding_idx=0),
+        SequenceFeature(name='user_hist_categories', vocab_size=category_vocab_size,
+                       max_len=sequence_max_len, embedding_dim=16, padding_idx=0),
+    ]
+    
+    # Item dense features
+    item_dense_features = [DenseFeature(name='item_price', input_dim=1)]
+    
+    # Item sparse features
+    item_sparse_features = [
+        SparseFeature(name='item_id', vocab_size=item_vocab_size, embedding_dim=item_embedding_dim),
+        SparseFeature(name='item_category', vocab_size=category_vocab_size, embedding_dim=16),
+        SparseFeature(name='item_brand', vocab_size=brand_vocab_size, embedding_dim=16),
+    ]
+    item_sparse_features.extend([
+        SparseFeature(name=f'item_feature_{i}', vocab_size=item_feature_vocab_size, embedding_dim=8)
+        for i in range(3)
+    ])
+    
+    # Item sequence features (empty for most match models)
+    item_sequence_features = []
+    
+    return (df, user_dense_features, user_sparse_features, user_sequence_features,
+            item_dense_features, item_sparse_features, item_sequence_features)
+
+
+def generate_multitask_data(
+    n_samples: int = 10000,
+    n_dense: int = 5,
+    n_sparse: int = 8,
+    n_sequences: int = 2,
+    user_vocab_size: int = 1000,
+    item_vocab_size: int = 500,
+    sparse_vocab_size: int = 50,
+    sequence_max_len: int = 20,
+    embedding_dim: int = 16,
+    seed: int = 42
+) -> Tuple[pd.DataFrame, List, List, List]:
+    """
+    Generate synthetic data for multi-task learning
+    
+    Returns:
+        tuple: (dataframe, dense_features, sparse_features, sequence_features)
+    """
+    print(f"Generating {n_samples} synthetic multi-task samples...")
+    
+    np.random.seed(seed)
+    data = {}
+    
+    # Generate dense features
+    for i in range(n_dense):
+        data[f'dense_{i}'] = np.random.randn(n_samples).astype(np.float32)
+    
+    # Generate sparse features
+    data['user_id'] = np.random.randint(1, user_vocab_size, n_samples)
+    data['item_id'] = np.random.randint(1, item_vocab_size, n_samples)
+    
+    for i in range(n_sparse - 2):
+        data[f'sparse_{i}'] = np.random.randint(1, sparse_vocab_size, n_samples)
+    
+    # Generate sequence features
+    sequence_names = []
+    sequence_vocabs = []
+    
+    for i in range(n_sequences):
+        sequences = []
+        for _ in range(n_samples):
+            seq_len = np.random.randint(5, sequence_max_len + 1)
+            if i == 0:
+                seq = np.random.randint(0, item_vocab_size, seq_len).tolist()
+                seq_vocab = item_vocab_size
+                seq_name = 'sequence_0'
+            else:
+                seq = np.random.randint(0, sparse_vocab_size, seq_len).tolist()
+                seq_vocab = sparse_vocab_size
+                seq_name = f'sequence_{i}'
+            
+            seq = seq + [0] * (sequence_max_len - len(seq))
+            sequences.append(seq)
+        
+        data[seq_name] = sequences
+        sequence_names.append(seq_name)
+        sequence_vocabs.append(seq_vocab)
+    
+    # Generate multi-task labels with correlation
+    # CTR (click) is relatively easier to predict
+    ctr_logits = (
+        data['dense_0'] * 0.3 + 
+        data['dense_1'] * 0.2 + 
+        np.random.randn(n_samples) * 0.5
+    )
+    data['click'] = (1 / (1 + np.exp(-ctr_logits)) > 0.5).astype(np.float32)
+    
+    # CVR (conversion) depends on click and is harder
+    cvr_logits = (
+        data['dense_2'] * 0.2 + 
+        data['dense_3'] * 0.15 +
+        data['click'] * 1.5 +  # Strong dependency on click
+        np.random.randn(n_samples) * 0.8
+    )
+    data['conversion'] = (1 / (1 + np.exp(-cvr_logits)) > 0.3).astype(np.float32)
+    
+    # CTCVR = click AND conversion
+    data['ctcvr'] = (data['click'] * data['conversion']).astype(np.float32)
+    
+    df = pd.DataFrame(data)
+    print(f"Generated data shape: {df.shape}")
+    print(f"Click rate: {data['click'].mean():.4f}")
+    print(f"Conversion rate (overall): {data['conversion'].mean():.4f}")
+    if data['click'].sum() > 0:
+        print(f"Conversion rate (given click): {data['conversion'][data['click'] == 1].mean():.4f}")
+    print(f"CTCVR rate: {data['ctcvr'].mean():.4f}")
+    
+    # Import here to avoid circular import
+    from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
+    
+    # Create feature definitions
+    dense_features = [DenseFeature(name=f'dense_{i}', input_dim=1) for i in range(n_dense)]
+    
+    # Create sparse features
+    sparse_features = [
+        SparseFeature(name='user_id', embedding_name='user_emb', 
+                     vocab_size=user_vocab_size, embedding_dim=embedding_dim),
+        SparseFeature(name='item_id', embedding_name='item_emb', 
+                     vocab_size=item_vocab_size, embedding_dim=embedding_dim),
+    ]
+    sparse_features.extend([
+        SparseFeature(name=f'sparse_{i}', embedding_name=f'sparse_{i}_emb',
+                     vocab_size=sparse_vocab_size, embedding_dim=embedding_dim) 
+        for i in range(n_sparse - 2)
+    ])
+    
+    # Create sequence features
+    sequence_features = []
+    for i, (seq_name, seq_vocab) in enumerate(zip(sequence_names, sequence_vocabs)):
+        if i == 0:
+            embedding_name = 'item_emb'
+        else:
+            embedding_name = 'sparse_0_emb'
+        sequence_features.append(
+            SequenceFeature(name=seq_name, vocab_size=seq_vocab, max_len=sequence_max_len, 
+                          embedding_dim=embedding_dim, padding_idx=0, embedding_name=embedding_name)
+        )
+    
+    return df, dense_features, sparse_features, sequence_features
+
+
+def generate_distributed_ranking_data(
+    num_samples: int = 100000,
+    num_users: int = 10000,
+    num_items: int = 5000,
+    num_categories: int = 20,
+    num_cities: int = 100,
+    max_seq_len: int = 50,
+    embedding_dim: int = 32,
+    seed: int = 42,
+) -> Tuple[pd.DataFrame, List, List, List]:
+    """
+    Generate synthetic data for distributed training scenarios
+    
+    Returns:
+        tuple: (dataframe, dense_features, sparse_features, sequence_features)
+    """
+    return generate_ranking_data(
+        n_samples=num_samples,
+        n_dense=5,
+        n_sparse=6,  # user_id, item_id + 4 custom features
+        n_sequences=2,
+        user_vocab_size=num_users + 1,
+        item_vocab_size=num_items + 1,
+        sequence_max_len=max_seq_len,
+        embedding_dim=embedding_dim,
+        seed=seed,
+        custom_sparse_features={'gender': 2, 'age_group': 7, 'category': num_categories,'city': num_cities},
+        use_simple_names=False
+    )