nextrec 0.3.3__tar.gz → 0.3.5__tar.gz

{nextrec-0.3.3 → nextrec-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.3.3
+Version: 0.3.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
@@ -63,7 +63,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.3.3-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.3.5-orange.svg)
 
 English | [中文文档](README_zh.md)
 
@@ -110,7 +110,7 @@ To dive deeper, Jupyter notebooks are available:
 - [Hands on the NextRec framework](/tutorials/notebooks/en/Hands%20on%20nextrec.ipynb)
 - [Using the data processor for preprocessing](/tutorials/notebooks/en/Hands%20on%20dataprocessor.ipynb)
 
-> Current version [0.3.3]: the matching module is not fully polished yet and may have compatibility issues or unexpected errors. Please raise an issue if you run into problems.
+> Current version [0.3.5]: the matching module is not fully polished yet and may have compatibility issues or unexpected errors. Please raise an issue if you run into problems.
 
 ## 5-Minute Quick Start
 

{nextrec-0.3.3 → nextrec-0.3.5}/README.md

@@ -7,7 +7,7 @@
 ![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.3.3-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.3.5-orange.svg)
 
 English | [中文文档](README_zh.md)
 
@@ -54,7 +54,7 @@ To dive deeper, Jupyter notebooks are available:
 - [Hands on the NextRec framework](/tutorials/notebooks/en/Hands%20on%20nextrec.ipynb)
 - [Using the data processor for preprocessing](/tutorials/notebooks/en/Hands%20on%20dataprocessor.ipynb)
 
-> Current version [0.3.3]: the matching module is not fully polished yet and may have compatibility issues or unexpected errors. Please raise an issue if you run into problems.
+> Current version [0.3.5]: the matching module is not fully polished yet and may have compatibility issues or unexpected errors. Please raise an issue if you run into problems.
 
 ## 5-Minute Quick Start
 

{nextrec-0.3.3 → nextrec-0.3.5}/README_zh.md

@@ -7,7 +7,7 @@
 ![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.3.3-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.3.5-orange.svg)
 
 [English Version](README.md) | 中文文档
 
@@ -54,7 +54,7 @@ NextRec采用模块化、低耦合的工程设计，使得推荐系统从数据
 - [如何上手NextRec框架](/tutorials/notebooks/zh/Hands%20on%20nextrec.ipynb)
 - [如何使用数据处理器进行数据预处理](/tutorials/notebooks/zh/Hands%20on%20dataprocessor.ipynb)
 
-> 当前版本[0.3.3]，召回模型模块尚不完善，可能存在一些兼容性问题或意外报错，如果遇到问题，欢迎开发者在Issue区提出问题。
+> 当前版本[0.3.5]，召回模型模块尚不完善，可能存在一些兼容性问题或意外报错，如果遇到问题，欢迎开发者在Issue区提出问题。
 
 ## 5分钟快速上手
 

{nextrec-0.3.3 → nextrec-0.3.5}/docs/rtd/conf.py

@@ -11,7 +11,7 @@ sys.path.insert(0, str(PROJECT_ROOT / "nextrec"))
 project = "NextRec"
 copyright = "2025, Yang Zhou"
 author = "Yang Zhou"
-release = "0.3.3"
+release = "0.3.5"
 
 extensions = [
     "myst_parser",

nextrec-0.3.5/nextrec/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.3.5"

{nextrec-0.3.3 → nextrec-0.3.5}/nextrec/basic/features.py

@@ -7,7 +7,7 @@ Author: Yang Zhou, zyaztec@gmail.com
 """
 import torch
 from nextrec.utils.embedding import get_auto_embedding_dim
-from nextrec.utils.common import normalize_to_list
+from nextrec.utils.feature import normalize_to_list
 
 class BaseFeature(object):
     def __repr__(self):

{nextrec-0.3.3 → nextrec-0.3.5}/nextrec/basic/loggers.py

@@ -2,17 +2,19 @@
 NextRec Basic Loggers
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 03/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
-
 import os
 import re
 import sys
+import json
 import copy
 import logging
-from nextrec.basic.session import create_session
+import numbers
+from typing import Mapping, Any
+from nextrec.basic.session import create_session, Session
 
 ANSI_CODES = {
     'black': '\033[30m',
@@ -77,17 +79,12 @@ def colorize(text: str, color: str | None = None, bold: bool = False) -> str:
     """Apply ANSI color and bold formatting to the given text."""
     if not color and not bold:
         return text
-
     result = ""
-
     if bold:
         result += ANSI_BOLD
-
     if color and color in ANSI_CODES:
         result += ANSI_CODES[color]
-
     result += text + ANSI_RESET
-
     return result
 
 def setup_logger(session_id: str | os.PathLike | None = None):
@@ -126,3 +123,69 @@ def setup_logger(session_id: str | os.PathLike | None = None):
     logger.addHandler(console_handler)
     
     return logger
+
+class TrainingLogger:
+    def __init__(
+        self,
+        session: Session,
+        enable_tensorboard: bool,
+        log_name: str = "training_metrics.jsonl",
+    ) -> None:
+        self.session = session
+        self.enable_tensorboard = enable_tensorboard
+        self.log_path = session.metrics_dir / log_name
+        self.log_path.parent.mkdir(parents=True, exist_ok=True)
+
+        self.tb_writer = None
+        self.tb_dir = None
+
+        if self.enable_tensorboard:
+            self._init_tensorboard()
+
+    def _init_tensorboard(self) -> None:
+        try:
+            from torch.utils.tensorboard import SummaryWriter  # type: ignore
+        except ImportError:
+            logging.warning("[TrainingLogger] tensorboard not installed, disable tensorboard logging.")
+            self.enable_tensorboard = False
+            return
+        tb_dir = self.session.logs_dir / "tensorboard"
+        tb_dir.mkdir(parents=True, exist_ok=True)
+        self.tb_dir = tb_dir
+        self.tb_writer = SummaryWriter(log_dir=str(tb_dir))
+
+    @property
+    def tensorboard_logdir(self):
+        return self.tb_dir
+
+    def format_metrics(self, metrics: Mapping[str, Any], split: str) -> dict[str, float]:
+        formatted: dict[str, float] = {}
+        for key, value in metrics.items():
+            if isinstance(value, numbers.Number):
+                formatted[f"{split}/{key}"] = float(value)
+            elif hasattr(value, "item"):
+                try:
+                    formatted[f"{split}/{key}"] = float(value.item())
+                except Exception:
+                    continue
+        return formatted
+
+    def log_metrics(self, metrics: Mapping[str, Any], step: int, split: str = "train") -> None:
+        payload = self.format_metrics(metrics, split)
+        payload["step"] = int(step)
+        with self.log_path.open("a", encoding="utf-8") as f:
+            f.write(json.dumps(payload, ensure_ascii=False) + "\n")
+
+        if not self.tb_writer:
+            return
+        step = int(payload.get("step", 0))
+        for key, value in payload.items():
+            if key == "step":
+                continue
+            self.tb_writer.add_scalar(key, value, global_step=step)
+
+    def close(self) -> None:
+        if self.tb_writer:
+            self.tb_writer.flush()
+            self.tb_writer.close()
+            self.tb_writer = None

{nextrec-0.3.3 → nextrec-0.3.5}/nextrec/basic/model.py

@@ -10,6 +10,8 @@ import os
 import tqdm
 import pickle
 import logging
+import getpass
+import socket
 import numpy as np
 import pandas as pd
 import torch
@@ -24,15 +26,17 @@ from nextrec.basic.callback import EarlyStopper
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature, FeatureSet
 from nextrec.data.dataloader import TensorDictDataset, RecDataLoader
 
-from nextrec.basic.loggers import setup_logger, colorize
+from nextrec.basic.loggers import setup_logger, colorize, TrainingLogger
 from nextrec.basic.session import resolve_save_path, create_session
 from nextrec.basic.metrics import configure_metrics, evaluate_metrics, check_user_id
 
 from nextrec.data.dataloader import build_tensors_from_data
-from nextrec.data.data_utils import get_column_data, collate_fn, batch_to_dict, get_user_ids
+from nextrec.data.data_processing import get_column_data, get_user_ids
+from nextrec.data.batch_utils import collate_fn, batch_to_dict
 
 from nextrec.loss import get_loss_fn, get_loss_kwargs
-from nextrec.utils import get_optimizer, get_scheduler, to_tensor
+from nextrec.utils import get_optimizer, get_scheduler
+from nextrec.utils.tensor import to_tensor
 
 from nextrec import __version__
 
@@ -88,6 +92,7 @@ class BaseModel(FeatureSet, nn.Module):
         self.early_stop_patience = early_stop_patience
         self.max_gradient_norm = 1.0   
         self.logger_initialized = False
+        self.training_logger: TrainingLogger | None = None
 
     def register_regularization_weights(self, embedding_attr: str = "embedding", exclude_modules: list[str] | None = None, include_modules: list[str] | None = None) -> None:
         exclude_modules = exclude_modules or []
@@ -275,11 +280,13 @@ class BaseModel(FeatureSet, nn.Module):
             metrics: list[str] | dict[str, list[str]] | None = None, # ['auc', 'logloss'] or {'target1': ['auc', 'logloss'], 'target2': ['mse']}
             epochs:int=1, shuffle:bool=True, batch_size:int=32,
             user_id_column: str | None = None,
-            validation_split: float | None = None):
+            validation_split: float | None = None,
+            tensorboard: bool = True,):
         self.to(self.device)
         if not self.logger_initialized:
             setup_logger(session_id=self.session_id)
             self.logger_initialized = True
+        self.training_logger = TrainingLogger(session=self.session, enable_tensorboard=tensorboard)
 
         self.metrics, self.task_specific_metrics, self.best_metrics_mode = configure_metrics(task=self.task, metrics=metrics, target_names=self.target_columns) # ['auc', 'logloss'], {'target1': ['auc', 'logloss'], 'target2': ['mse']}, 'max'
         self.early_stopper = EarlyStopper(patience=self.early_stop_patience, mode=self.best_metrics_mode)
@@ -303,6 +310,20 @@ class BaseModel(FeatureSet, nn.Module):
             is_streaming = True
 
         self.summary()
+        logging.info("")
+        if self.training_logger and self.training_logger.enable_tensorboard:
+            tb_dir = self.training_logger.tensorboard_logdir
+            if tb_dir:
+                user = getpass.getuser()
+                host = socket.gethostname()
+                tb_cmd = f"tensorboard --logdir {tb_dir} --port 6006"
+                ssh_hint = f"ssh -L 6006:localhost:6006 {user}@{host}"
+                logging.info(colorize(f"TensorBoard logs saved to: {tb_dir}", color="cyan"))
+                logging.info(colorize("To view logs, run:", color="cyan"))
+                logging.info(colorize(f"    {tb_cmd}", color="cyan"))
+                logging.info(colorize("Then SSH port forward:", color="cyan"))
+                logging.info(colorize(f"    {ssh_hint}", color="cyan"))
+
         logging.info("")
         logging.info(colorize("=" * 80, bold=True))
         if is_streaming:
@@ -312,7 +333,7 @@ class BaseModel(FeatureSet, nn.Module):
         logging.info(colorize("=" * 80, bold=True))
         logging.info("")
         logging.info(colorize(f"Model device: {self.device}", bold=True))
-    
+
         for epoch in range(epochs):
             self.epoch_index = epoch
             if is_streaming:
@@ -326,7 +347,8 @@ class BaseModel(FeatureSet, nn.Module):
             else:
                 train_loss = train_result
                 train_metrics = None
-            
+
+            train_log_payload: dict[str, float] = {}
             # handle logging for single-task and multi-task
             if self.nums_task == 1:
                 log_str = f"Epoch {epoch + 1}/{epochs} - Train: loss={train_loss:.4f}"
@@ -334,6 +356,9 @@ class BaseModel(FeatureSet, nn.Module):
                     metrics_str = ", ".join([f"{k}={v:.4f}" for k, v in train_metrics.items()])
                     log_str += f", {metrics_str}"
                 logging.info(colorize(log_str))
+                train_log_payload["loss"] = float(train_loss)
+                if train_metrics:
+                    train_log_payload.update(train_metrics)
             else:
                 total_loss_val = np.sum(train_loss) if isinstance(train_loss, np.ndarray) else train_loss  # type: ignore
                 log_str = f"Epoch {epoch + 1}/{epochs} - Train: loss={total_loss_val:.4f}"
@@ -356,12 +381,17 @@ class BaseModel(FeatureSet, nn.Module):
                                 task_metric_strs.append(f"{target_name}[{metrics_str}]")
                         log_str += ", " + ", ".join(task_metric_strs)
                 logging.info(colorize(log_str))
+                train_log_payload["loss"] = float(total_loss_val)
+                if train_metrics:
+                    train_log_payload.update(train_metrics)
+            if self.training_logger:
+                self.training_logger.log_metrics(train_log_payload, step=epoch + 1, split="train")
             if valid_loader is not None:
                 # pass user_ids only if needed for GAUC metric
                 val_metrics = self.evaluate(valid_loader, user_ids=valid_user_ids if self.needs_user_ids else None) # {'auc': 0.75, 'logloss': 0.45} or {'auc_target1': 0.75, 'logloss_target1': 0.45, 'mse_target2': 3.2}
                 if self.nums_task == 1:
                     metrics_str = ", ".join([f"{k}={v:.4f}" for k, v in val_metrics.items()])
-                    logging.info(colorize(f"Epoch {epoch + 1}/{epochs} - Valid: {metrics_str}", color="cyan"))
+                    logging.info(colorize(f"  Epoch {epoch + 1}/{epochs} - Valid: {metrics_str}", color="cyan"))
                 else:
                     # multi task metrics
                     task_metrics = {}
@@ -378,7 +408,9 @@ class BaseModel(FeatureSet, nn.Module):
                         if target_name in task_metrics:
                             metrics_str = ", ".join([f"{k}={v:.4f}" for k, v in task_metrics[target_name].items()])
                             task_metric_strs.append(f"{target_name}[{metrics_str}]")
-                    logging.info(colorize(f"Epoch {epoch + 1}/{epochs} - Valid: " + ", ".join(task_metric_strs), color="cyan"))
+                    logging.info(colorize(f"  Epoch {epoch + 1}/{epochs} - Valid: " + ", ".join(task_metric_strs), color="cyan"))
+                if val_metrics and self.training_logger:
+                    self.training_logger.log_metrics(val_metrics, step=epoch + 1, split="valid")
                 # Handle empty validation metrics
                 if not val_metrics:
                     self.save_model(self.checkpoint_path, add_timestamp=False, verbose=False)
@@ -401,6 +433,7 @@ class BaseModel(FeatureSet, nn.Module):
                         self.best_metric = primary_metric
                         improved = True
                 self.save_model(self.checkpoint_path, add_timestamp=False, verbose=False)
+                logging.info(" ")
                 if improved:
                     logging.info(colorize(f"Validation {primary_metric_key} improved to {self.best_metric:.4f}"))
                     self.save_model(self.best_path, add_timestamp=False, verbose=False)
@@ -431,6 +464,8 @@ class BaseModel(FeatureSet, nn.Module):
         if valid_loader is not None:
             logging.info(colorize(f"Load best model from: {self.best_checkpoint_path}"))
             self.load_model(self.best_checkpoint_path, map_location=self.device, verbose=False)
+        if self.training_logger:
+            self.training_logger.close()
         return self
 
     def train_epoch(self, train_loader: DataLoader, is_streaming: bool = False) -> Union[float, np.ndarray, tuple[Union[float, np.ndarray], dict]]:
@@ -527,6 +562,7 @@ class BaseModel(FeatureSet, nn.Module):
                     batch_user_id = get_user_ids(data=batch_dict, id_columns=self.id_columns)
                     if batch_user_id is not None:
                         collected_user_ids.append(batch_user_id)
+        logging.info(" ")
         logging.info(colorize(f"  Evaluation batches processed: {batch_count}", color="cyan"))
         if len(y_true_list) > 0:
             y_true_all = np.concatenate(y_true_list, axis=0)
@@ -956,9 +992,7 @@ class BaseModel(FeatureSet, nn.Module):
         logger.info(f"  Session ID:            {self.session_id}")
         logger.info(f"  Features Config Path:  {self.features_config_path}")
         logger.info(f"  Latest Checkpoint:     {self.checkpoint_path}")
-        
-        logger.info("")
-        logger.info("")
+
 
 
 class BaseMatchModel(BaseModel):

{nextrec-0.3.3 → nextrec-0.3.5}/nextrec/basic/session.py

@@ -1,14 +1,5 @@
 """Session and experiment utilities.
 
-This module centralizes session/experiment management so the rest of the
-framework writes all artifacts to a consistent location:: <pwd>/log/<experiment_id>/
-
-Within that folder we keep model parameters, checkpoints, training metrics,
-evaluation metrics, and consolidated log output. When users do not provide an
-``experiment_id`` a timestamp-based identifier is generated once per process to
-avoid scattering files across multiple directories. Test runs are redirected to
-temporary folders so local trees are not polluted.
-
 Date: create on 23/11/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
@@ -16,7 +7,7 @@ Author: Yang Zhou,zyaztec@gmail.com
 import os
 import tempfile
 from dataclasses import dataclass
-from datetime import datetime
+from datetime import datetime, timezone
 from pathlib import Path
 
 __all__ = [
@@ -74,6 +65,7 @@ def create_session(experiment_id: str | Path | None = None) -> Session:
     if experiment_id is not None and str(experiment_id).strip():
         exp_id = str(experiment_id).strip()
     else:
+        # Use local time for session naming
         exp_id = "nextrec_session_" + datetime.now().strftime("%Y%m%d")
 
     if (
@@ -111,6 +103,7 @@ def resolve_save_path(
       timestamp.
     - Parent directories are created.
     """
+    # Use local time for file timestamps
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") if add_timestamp else None
 
     normalized_suffix = suffix if suffix.startswith(".") else f".{suffix}"

{nextrec-0.3.3 → nextrec-0.3.5}/nextrec/data/__init__.py

@@ -1,48 +1,86 @@
 """
 Data utilities package for NextRec
 
-This package provides data processing and manipulation utilities.
+This package provides data processing and manipulation utilities organized by category:
+- batch_utils: Batch collation and processing
+- data_processing: Data manipulation and user ID extraction
+- data_utils: Legacy module (re-exports from specialized modules)
+- dataloader: Dataset and DataLoader implementations
+- preprocessor: Data preprocessing pipeline
 
 Date: create on 13/11/2025
+Last update: 03/12/2025 (refactored)
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
-from nextrec.data.data_utils import (
-    collate_fn,
+# Batch utilities
+from nextrec.data.batch_utils import collate_fn, batch_to_dict, stack_section
+
+# Data processing utilities
+from nextrec.data.data_processing import (
     get_column_data,
-    default_output_dir,
     split_dict_random,
     build_eval_candidates,
+    get_user_ids,
+)
+
+# File utilities (from utils package)
+from nextrec.utils.file import (
     resolve_file_paths,
     iter_file_chunks,
     read_table,
     load_dataframes,
+    default_output_dir,
 )
-from nextrec.basic.features import FeatureSet
-from nextrec.data import data_utils
+
+# DataLoader components
 from nextrec.data.dataloader import (
     TensorDictDataset,
     FileDataset,
     RecDataLoader,
     build_tensors_from_data,
 )
+
+# Preprocessor
 from nextrec.data.preprocessor import DataProcessor
 
+# Feature definitions
+from nextrec.basic.features import FeatureSet
+
+# Legacy module (for backward compatibility)
+from nextrec.data import data_utils
+
 __all__ = [
+    # Batch utilities
     'collate_fn',
+    'batch_to_dict',
+    'stack_section',
+    
+    # Data processing
     'get_column_data',
-    'default_output_dir',
     'split_dict_random',
     'build_eval_candidates',
+    'get_user_ids',
+    
+    # File utilities
     'resolve_file_paths',
     'iter_file_chunks',
     'read_table',
     'load_dataframes',
-    'FeatureSet',
-    'data_utils',
+    'default_output_dir',
+    
+    # DataLoader
     'TensorDictDataset',
     'FileDataset',
     'RecDataLoader',
     'build_tensors_from_data',
+    
+    # Preprocessor
     'DataProcessor',
+    
+    # Features
+    'FeatureSet',
+    
+    # Legacy module
+    'data_utils',
 ]

nextrec-0.3.5/nextrec/data/batch_utils.py

@@ -0,0 +1,80 @@
+"""
+Batch collation utilities for NextRec
+
+Date: create on 03/12/2025
+Author: Yang Zhou, zyaztec@gmail.com
+"""
+
+import torch
+import numpy as np
+from typing import Any, Mapping
+
+def stack_section(batch: list[dict], section: str):
+    entries = [item.get(section) for item in batch if item.get(section) is not None]
+    if not entries:
+        return None
+    merged: dict = {}
+    for name in entries[0]:  # type: ignore
+        tensors = [item[section][name] for item in batch if item.get(section) is not None and name in item[section]]
+        merged[name] = torch.stack(tensors, dim=0)
+    return merged
+
+def collate_fn(batch):
+    """
+    Collate a list of sample dicts into the unified batch format:
+    {
+        "features": {name: Tensor(B, ...)},
+        "labels": {target: Tensor(B, ...)} or None,
+        "ids": {id_name: Tensor(B, ...)} or None,
+    }
+    Args: batch: List of samples from DataLoader
+        
+    Returns: dict: Batched data in unified format
+    """
+    if not batch:
+        return {"features": {}, "labels": None, "ids": None}
+
+    first = batch[0]
+    if isinstance(first, dict) and "features" in first:
+        # Streaming dataset yields already-batched chunks; avoid adding an extra dim.
+        if first.get("_already_batched") and len(batch) == 1:
+            return {
+                "features": first.get("features", {}),
+                "labels": first.get("labels"),
+                "ids": first.get("ids"),
+            }
+        return {
+            "features": stack_section(batch, "features") or {},
+            "labels": stack_section(batch, "labels"),
+            "ids": stack_section(batch, "ids"),
+        }
+
+    # Fallback: stack tuples/lists of tensors
+    num_tensors = len(first)
+    result = []
+    for i in range(num_tensors):
+        tensor_list = [item[i] for item in batch]
+        first_item = tensor_list[0]
+        if isinstance(first_item, torch.Tensor):
+            stacked = torch.cat(tensor_list, dim=0)
+        elif isinstance(first_item, np.ndarray):
+            stacked = np.concatenate(tensor_list, axis=0)
+        elif isinstance(first_item, list):
+            combined = []
+            for entry in tensor_list:
+                combined.extend(entry)
+            stacked = combined
+        else:
+            stacked = tensor_list
+        result.append(stacked)
+    return tuple(result)
+
+
+def batch_to_dict(batch_data: Any, include_ids: bool = True) -> dict:
+    if not (isinstance(batch_data, Mapping) and "features" in batch_data):
+        raise TypeError("[BaseModel-batch_to_dict Error] Batch data must be a dict with 'features' produced by the current DataLoader.")
+    return {
+        "features": batch_data.get("features", {}),
+        "labels": batch_data.get("labels"),
+        "ids": batch_data.get("ids") if include_ids else None,
+    }