style: apply project code formatter

Author: TyndaleLym

examples/generative/data/ml-1m/preprocess_ml_hstu.py

@@ -280,9 +280,24 @@ def generate_leave_last_out_samples(self, user_sequences, user_timestamps, vocab
         item_to_idx = vocab['item_to_idx'] if 'item_to_idx' in vocab else vocab
 
         splits = {
-            'train': {'seq_tokens': [], 'seq_positions': [], 'seq_time_diffs': [], 'targets': []},
-            'val': {'seq_tokens': [], 'seq_positions': [], 'seq_time_diffs': [], 'targets': []},
-            'test': {'seq_tokens': [], 'seq_positions': [], 'seq_time_diffs': [], 'targets': []},
+            'train': {
+                'seq_tokens': [],
+                'seq_positions': [],
+                'seq_time_diffs': [],
+                'targets': []
+            },
+            'val': {
+                'seq_tokens': [],
+                'seq_positions': [],
+                'seq_time_diffs': [],
+                'targets': []
+            },
+            'test': {
+                'seq_tokens': [],
+                'seq_positions': [],
+                'seq_time_diffs': [],
+                'targets': []
+            },
         }
 
         positions_template = list(range(self.max_seq_len))

examples/generative/run_hstu_movielens.py

@@ -235,10 +235,12 @@ def main(
         optimizer_params={
             "lr": learning_rate,
             "weight_decay": weight_decay,
-            "betas": (0.9, 0.98),
+            "betas": (0.9,
+                      0.98),
         },
         n_epoch=epoch,
-        earlystop_patience=max(epoch + 1, 10),
+        earlystop_patience=max(epoch + 1,
+                               10),
         device=device,
         model_path=save_dir,
     )

tests/test_hstu_model.py

@@ -174,9 +174,13 @@ def test_vocab_mask_drops_per_user_history():
 
 def test_seqtrainer_next_token_loss_masks_left_pad_positions():
     trainer = SeqTrainer(
-        nn.Linear(1, 1),
+        nn.Linear(1,
+                  1),
         device='cpu',
-        loss_params={"ignore_index": 0, "reduction": "sum"},
+        loss_params={
+            "ignore_index": 0,
+            "reduction": "sum"
+        },
     )
     logits = torch.zeros(1, 4, 10)
     seq_tokens = torch.tensor([[0, 0, 5, 6]])

torch_rechub/basic/layers.py

@@ -948,8 +948,7 @@ def __init__(self, d_model=512, n_heads=8, n_layers=4, dqk=64, dv=64, dropout=0.
                     time_bucket_fn=time_bucket_fn,
                     time_bucket_divisor=time_bucket_divisor,
                     time_bucket_unit=time_bucket_unit,
-                )
-                for _ in range(n_layers)
+                ) for _ in range(n_layers)
             ]
         )
 

torch_rechub/models/generative/hstu.py

@@ -227,7 +227,8 @@ def forward(self, x, time_diffs=None):
         """
         batch_size, seq_len = x.shape
         if seq_len > self.max_seq_len:
-            raise ValueError(f"Input seq_len ({seq_len}) exceeds max_seq_len ({self.max_seq_len}). " f"Either truncate the input or rebuild the model with a larger max_seq_len.")
+            raise ValueError(f"Input seq_len ({seq_len}) exceeds max_seq_len ({self.max_seq_len}). "
+                             f"Either truncate the input or rebuild the model with a larger max_seq_len.")
 
         padding_mask = x.ne(0)  # (B, L) — True for valid tokens
 

torch_rechub/trainers/seq_trainer.py

@@ -348,7 +348,8 @@ def export_onnx(self, output_path, batch_size=2, seq_length=50, vocab_size=None,
                     overlap = set(export_kwargs.keys()) & set(onnx_export_kwargs.keys())
                     overlap.discard("dynamo")
                     if overlap:
-                        raise ValueError("onnx_export_kwargs contains keys that overlap with explicit args: " f"{sorted(overlap)}. Please set them via export_onnx() parameters instead.")
+                        raise ValueError("onnx_export_kwargs contains keys that overlap with explicit args: "
+                                         f"{sorted(overlap)}. Please set them via export_onnx() parameters instead.")
                     export_kwargs.update(onnx_export_kwargs)
 
                 # Auto-pick exporter:
@@ -473,7 +474,8 @@ def visualization(self, seq_length=50, vocab_size=None, batch_size=2, depth=3, s
             elif hasattr(model, 'item_num'):
                 vocab_size = model.item_num
             else:
-                raise ValueError("vocab_size must be provided or model must have " "'vocab_size' or 'item_num' attribute")
+                raise ValueError("vocab_size must be provided or model must have "
+                                 "'vocab_size' or 'item_num' attribute")
 
         # Generate dummy inputs for sequence model
         dummy_seq_tokens = torch.randint(0, vocab_size, (batch_size, seq_length), device=viz_device)