Fix code discrepancies encountered during the open-sourcing process

.gitignore

@@ -22,3 +22,8 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+
+cache/
+logs/
+tensorboard/
+ckpt/

README.md

@@ -1,7 +1,7 @@
 # ActionPiece: Contextual Action Tokenization
 
 This repository provides the code for implementing ActionPiece described in our
-**ICML 25 Spotlight** paper "[Contextually Tokenizing Action Sequences forGenerative Recommendation](https://arxiv.org/abs/2502.13581)".
+**ICML 25 Spotlight** paper "[Contextually Tokenizing Action Sequences for Generative Recommendation](https://arxiv.org/abs/2502.13581)".
 
 Unlike existing generative recommendation (GR) models that tokenize each action
 independently, we propose ActionPiece, a method that

genrec/dataset.py

@@ -19,7 +19,6 @@
 
 from typing import Any
 import datasets as datasets_lib
-from genrec.utils import log as log_lib
 
 
 class AbstractDataset:
@@ -157,6 +156,7 @@ def split(self) -> dict[str, datasets_lib.Dataset]:
 
   def log(self, message: str, level: str = 'info') -> None:
     """Logs a message with the specified level."""
+    from genrec.utils import log as log_lib
     return log_lib(
         message, self.config['accelerator'], self.logger, level=level
     )

genrec/datasets/AmazonReviews2014/dataset.py

@@ -15,6 +15,7 @@
 
 """Dataset for Amazon Reviews 2014."""
 
+import ast
 import collections
 import gzip
 import json
@@ -80,8 +81,18 @@ def parse_gz(path: str):
   """
   with gzip.open(path, 'r') as g:
     for l in g:
-      l = l.replace(b'true', b'True').replace(b'false', b'False')
-      yield json.loads(l)
+      try:
+        # Try to parse as standard JSON first
+        yield json.loads(l.decode('utf-8'))
+      except json.JSONDecodeError:
+        try:
+          # If that fails, try with Python literal evaluation after some replacements
+          l_str = l.decode('utf-8')
+          l_str = l_str.replace('true', 'True').replace('false', 'False').replace('null', 'None')
+          yield ast.literal_eval(l_str)
+        except (ValueError, SyntaxError):
+          # If both fail, skip this line and continue
+          continue
 
 
 def get_item_seqs(
@@ -274,7 +285,7 @@ def _load_metadata(
     self.log('[DATASET] Loading metadata...')
     data = {}
     item_asins = set(item2id.keys())
-    for info in tqdm.tqdm(self._parse_gz(path)):
+    for info in tqdm.tqdm(parse_gz(path)):
       if info['asin'] not in item_asins:
         continue
       data[info['asin']] = info
@@ -319,7 +330,7 @@ def _extract_meta_sentences(self, metadata: dict[str, Any]) -> dict[str, str]:
     """
     self.log('[DATASET] Extracting meta sentences...')
     item2meta = {}
-    for item, meta in tqdm(metadata.items()):
+    for item, meta in tqdm.tqdm(metadata.items()):
       meta_sentence = ''
       keys = set(meta.keys())
       features_needed = [

genrec/default.yaml

@@ -23,15 +23,15 @@ rand_seed: 2024
 reproducibility: True
 
 train_batch_size: 256
-eval_batch_size: 32
+eval_batch_size: 128
 lr: 0.003
 weight_decay: 0.1
 warmup_steps: 10000
 steps: ~
 epochs: 200
 max_grad_norm: 1.0      # None for no clipping, else a float value
 eval_interval: 1        # Evaluate every n epochs
-patience: 20            # Early stopping. Stop training after n epochs without improvement. Set to None to disable
+patience: 50            # Early stopping. Stop training after n epochs without improvement. Set to None to disable
 
 topk: [5,10,20,50]
 metrics: [ndcg,recall,err]

genrec/models/ActionPiece/config.yaml

@@ -17,7 +17,7 @@
 sent_emb_model: sentence-transformers/sentence-t5-base
 sent_emb_batch_size: 512
 sent_emb_dim: 768
-sent_emb_pca: 128 # -1 means no PCA, otherwise PCA dimension
+sent_emb_pca: -1 # -1 means no PCA, otherwise PCA dimension
 
 # Config for features
 n_threads: 32
@@ -29,7 +29,6 @@ n_hash_buckets: 128
 actionpiece_vocab_size: 40000
 
 # Config for the model
-n_prob_encode_plus: 0
 num_beams: 50 # Number of beams for beam search
 n_inference_ensemble: 5 # Number of inference ensemble
 train_shuffle: feature # none / feature / token

genrec/models/ActionPiece/model.py

@@ -102,8 +102,9 @@ def generate(self, batch: dict[Any, Any], n_return_sequences: int = 1):
     """
     n_ensemble = 1
     if self.n_inference_ensemble != -1:
-      assert batch['input_ids'].shape[0] % self.n_inference_ensemble == 0
-      n_ensemble = self.n_inference_ensemble
+      if batch['input_ids'].shape[0] != batch['labels'].shape[0]:
+        assert batch['input_ids'].shape[0] % self.n_inference_ensemble == 0
+        n_ensemble = self.n_inference_ensemble
     batch_size = batch['input_ids'].shape[0] // n_ensemble
 
     outputs = self.beam_search(

genrec/models/ActionPiece/tokenizer.py

@@ -41,7 +41,6 @@ class ActionPieceTokenizer(AbstractTokenizer):
       actionpiece (ActionPieceCore): ActionPiece core tokenizer.
       bos_token (int): The beginning token.
       eos_token (int): The end token.
-      n_prob_encode_plus (int): The number of probability encoding plus.
       n_inference_ensemble (int): The number of inference ensemble.
       train_shuffle (str): The shuffle strategy for training.
       encoded_labels (dict): A dictionary mapping label sequences to their
@@ -58,7 +57,6 @@ def __init__(self, config: dict[Any, Any], dataset: AbstractDataset):
     self.actionpiece = self._init_tokenizer(dataset)
     self.bos_token = self.actionpiece.vocab_size
     self.eos_token = self.actionpiece.vocab_size + 1
-    self.n_prob_encode_plus = self.config['n_prob_encode_plus']
     self.n_inference_ensemble = config['n_inference_ensemble']
     self.train_shuffle = config['train_shuffle']
     self.encoded_labels = {}
@@ -222,7 +220,7 @@ def _get_sem_ids(self, dataset: AbstractDataset) -> dict[Any, Any]:
       with open(sem_ids_path, 'r') as f:
         item2sem_ids = json.load(f)
       return {
-          k: v[: self.config['rq_n_codebooks']] for k, v in item2sem_ids.items()
+          k: v[: self.config['pq_n_codebooks']] for k, v in item2sem_ids.items()
       }
 
   def _get_attr_ids(self, dataset: AbstractDataset):

genrec/pipeline.py

@@ -17,7 +17,7 @@
 
 import logging
 import os
-from typing import Any
+from typing import Any, Dict, Union
 
 import accelerate as accelerate_lib
 from genrec import utils
@@ -59,11 +59,11 @@ class Pipeline:
 
   def __init__(
       self,
-      model_name: str | AbstractModel,
-      dataset_name: str | AbstractDataset,
-      tokenizer: AbstractTokenizer | None = None,
+      model_name: Union[str, AbstractModel],
+      dataset_name: Union[str, AbstractDataset],
+      tokenizer: Union[AbstractTokenizer, None] = None,
       trainer=None,
-      config_dict: dict[str, Any] | None = None,
+      config_dict: Union[Dict[str, Any], None] = None,
       config_file: str = None,
   ):
     self.config = utils.get_config(
@@ -151,15 +151,17 @@ def get_dataloader(split, batch_size, shuffle):
     train_dataloader = get_dataloader(
         'train', self.config['train_batch_size'], True
     )
+    val_dataloader = get_dataloader(
+        'val', self.config['eval_batch_size'], False
+    )
     if self.config['n_inference_ensemble'] == -1:
-      eval_batch_size = self.config['eval_batch_size']
+      test_batch_size = self.config['eval_batch_size']
     else:
-      eval_batch_size = max(
+      test_batch_size = max(
           self.config['eval_batch_size'] // self.config['n_inference_ensemble'],
           1,
       )
-    val_dataloader = get_dataloader('val', eval_batch_size, False)
-    test_dataloader = get_dataloader('test', eval_batch_size, False)
+    test_dataloader = get_dataloader('test', test_batch_size, False)
 
     self.trainer.fit(train_dataloader, val_dataloader)
 

genrec/utils.py

@@ -30,7 +30,6 @@
 import datasets.utils.logging
 from genrec.dataset import AbstractDataset
 from genrec.model import AbstractModel
-from genrec.trainer import Trainer
 import numpy as np
 import requests
 import torch
@@ -158,12 +157,21 @@ def log(message, accelerator, logger, level='info'):
       level (str): The log level ('info', 'error', 'warning', 'debug').
   """
   if accelerator.is_main_process:
+    # Map level names to their numeric values for compatibility with older Python versions
+    level_mapping = {
+        'DEBUG': logging.DEBUG,
+        'INFO': logging.INFO,
+        'WARNING': logging.WARNING,
+        'ERROR': logging.ERROR,
+        'CRITICAL': logging.CRITICAL
+    }
+
     try:
-      level = logging.getLevelNamesMapping()[level.upper()]
+      level_num = level_mapping[level.upper()]
     except KeyError as exc:
       raise ValueError(f'Invalid log level: {level}') from exc
 
-    logger.log(level, message)
+    logger.log(level_num, message)
 
 
 def get_tokenizer(model_name: str):
@@ -182,7 +190,7 @@ def get_tokenizer(model_name: str):
   module_name = f'genrec.models.{model_name}.tokenizer'
   try:
     module = importlib.import_module(module_name)
-    getattr(module, f'{model_name}Tokenizer')
+    return getattr(module, f'{model_name}Tokenizer')
   except Exception as exc:
     raise ValueError(f'Tokenizer for model "{model_name}" not found.') from exc
 
@@ -246,12 +254,16 @@ def get_trainer(model_name: Union[str, AbstractModel]):
       trainer_class: The trainer class corresponding to the given model name. If
       the model name is not found, the default Trainer class is returned.
   """
+  from genrec.trainer import Trainer
   if isinstance(model_name, str):
-    trainer_class = getattr(
-        importlib.import_module(f'genrec.models.{model_name}.trainer'),
-        f'{model_name}Trainer',
-    )
-    return trainer_class
+    try:
+      trainer_class = getattr(
+          importlib.import_module(f'genrec.models.{model_name}.trainer'),
+          f'{model_name}Trainer',
+      )
+      return trainer_class
+    except (ImportError, AttributeError):
+      return Trainer
 
   return Trainer
 
@@ -286,6 +298,18 @@ def _convert_value(value: str) -> Any:
     return True
   if value.lower() == 'false':
     return False
+
+  # Try to use eval for complex types (list, dict, tuple) but with safety checks
+  try:
+    new_v = eval(value)
+    if new_v is not None and isinstance(
+        new_v, (str, int, float, bool, list, dict, tuple)
+    ):
+      return new_v
+  except (NameError, SyntaxError, TypeError, ValueError):
+    pass
+
+  # Try basic numeric conversions
   try:
     return int(value)
   except ValueError:
@@ -294,10 +318,7 @@ def _convert_value(value: str) -> Any:
     return float(value)
   except ValueError:
     pass
-  try:
-    return list(map(lambda x: x.strip(), value.strip('[]').split(',')))
-  except (ValueError, TypeError):
-    pass
+
   return value