fix default args

webis-de · Nov 21, 2024 · 0164c53 · 0164c53
1 parent 0eaa265
commit 0164c53
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/lightning_ir/loss/loss.py b/lightning_ir/loss/loss.py
@@ -338,14 +338,12 @@ def get_ib_idcs(self, output: LightningIROutput, batch: TrainBatch) -> Tuple[tor
 
 class ScoreBasedInBatchLossFunction(InBatchLossFunction):
 
-    def __init__(
-        self,
-        min_target_diff: float,
-        pos_sampling_technique: Literal["first"] = "first",
-        neg_sampling_technique: Literal["all_and_non_first"] = "all_and_non_first",
-        max_num_neg_samples: int | None = None,
-    ):
-        super().__init__(pos_sampling_technique, neg_sampling_technique, max_num_neg_samples)
+    def __init__(self, min_target_diff: float, max_num_neg_samples: int | None = None):
+        super().__init__(
+            pos_sampling_technique="first",
+            neg_sampling_technique="all_and_non_first",
+            max_num_neg_samples=max_num_neg_samples,
+        )
         self.min_target_diff = min_target_diff
 
     def _sort_mask(self, mask: torch.Tensor, num_queries: int, num_docs: int, batch: TrainBatch) -> torch.Tensor: