whoops

terrierteam · Nov 23, 2024 · 96d2106 · 96d2106
1 parent aaa868c
commit 96d2106
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pyterrier_dr/biencoder.py b/pyterrier_dr/biencoder.py
@@ -72,7 +72,7 @@ def encode(self, texts, batch_size=None) -> np.array:
         return self.bi_encoder_model.encode_queries(texts, batch_size=batch_size or self.batch_size)
 
     def transform(self, inp: pd.DataFrame) -> pd.DataFrame:
-        pta.validate.columns(includes=['query'])
+        pta.validate.columns(inp, includes=['query'])
         it = inp['query'].values
         it, inv = np.unique(it, return_inverse=True)
         if self.verbose:
@@ -95,7 +95,7 @@ def encode(self, texts, batch_size=None) -> np.array:
         return self.bi_encoder_model.encode_docs(texts, batch_size=batch_size or self.batch_size)
 
     def transform(self, inp: pd.DataFrame) -> pd.DataFrame:
-        pta.validate.columns(includes=[self.text_field])
+        pta.validate.columns(inp, includes=[self.text_field])
         it = inp[self.text_field]
         if self.verbose:
             it = pt.tqdm(it, desc='Encoding Docs', unit='doc')