weaken tests

Author: teodor-r

rectools/models/nn/transformers/hstu.py

@@ -411,7 +411,8 @@ class HSTUModelConfig(TransformerModelConfig):
 
 class HSTUModel(TransformerModelBase[HSTUModelConfig]):
     """
-    HSTU model: transformer-based sequential model with unidirectional pointwise aggregated attention mechanism, combined with "Shifted Sequence" training objective.
+    HSTU model: transformer-based sequential model with unidirectional pointwise aggregated attention mechanism,
+    combined with "Shifted Sequence" training objective.
     Our implementation covers multiple loss functions and a variable number of negatives for them.
 
     References

tests/models/nn/transformers/test_sasrec.py

@@ -411,9 +411,9 @@ def test_u2i_losses(
         (
             pd.DataFrame(
                 {
-                    Columns.User: [10, 10, 10, 30, 30, 30, 40, 40, 40],
-                    Columns.Item: [13, 17, 12, 11, 13, 15, 17, 13, 11],
-                    Columns.Rank: [1, 2, 3, 1, 2, 3, 1, 2, 3],
+                    Columns.User: [30, 30, 30, 40, 40, 40],
+                    Columns.Item: [11, 13, 15, 17, 13, 11],
+                    Columns.Rank: [1, 2, 3, 1, 2, 3],
                 }
             ),
         ),
@@ -439,7 +439,7 @@ def test_u2i_with_key_and_attn_masks(
             similarity_module_type=DistanceSimilarityModule,
         )
         model.fit(dataset=dataset)
-        users = np.array([10, 30, 40])
+        users = np.unique(expected[Columns.User])
         actual = model.recommend(users=users, dataset=dataset, k=3, filter_viewed=False)
         pd.testing.assert_frame_equal(actual.drop(columns=Columns.Score), expected)
         pd.testing.assert_frame_equal(
@@ -452,9 +452,9 @@ def test_u2i_with_key_and_attn_masks(
         (
             pd.DataFrame(
                 {
-                    Columns.User: [10, 10, 10, 30, 30, 30, 40, 40, 40],
-                    Columns.Item: [13, 12, 11, 11, 12, 13, 13, 14, 12],
-                    Columns.Rank: [1, 2, 3, 1, 2, 3, 1, 2, 3],
+                    Columns.User: [30, 30, 30, 40, 40, 40],
+                    Columns.Item: [11, 12, 13, 13, 14, 12],
+                    Columns.Rank: [1, 2, 3, 1, 2, 3],
                 }
             ),
         ),
@@ -480,7 +480,7 @@ def test_u2i_with_item_features(
             similarity_module_type=DistanceSimilarityModule,
         )
         model.fit(dataset=dataset_item_features)
-        users = np.array([10, 30, 40])
+        users = np.unique(expected[Columns.User])
         actual = model.recommend(users=users, dataset=dataset_item_features, k=3, filter_viewed=False)
         pd.testing.assert_frame_equal(actual.drop(columns=Columns.Score), expected)
         pd.testing.assert_frame_equal(