Fixes

Author: CodingTil

py_css/models/base.py

@@ -192,15 +192,21 @@ def search(self, query: Query, context: Context) -> Tuple[Context, pd.DataFrame]
         if t5_rewriter_module.COPY_REWRITTEN_QUERY_COLUMN in result.columns:
             temp_result = result[result["qid"] == query.query_id]
             if not temp_result.empty:
-                query.query = temp_result.at[temp_result.index[0], t5_rewriter_module.COPY_REWRITTEN_QUERY_COLUMN]
+                query.query = temp_result[
+                    t5_rewriter_module.COPY_REWRITTEN_QUERY_COLUMN
+                ].iloc[0]
             else:
-                warnings.warn(f"Query {query.query_id} not found in result. This should not happen. All query-ids: {result['qid'].unique()}")
+                warnings.warn(
+                    f"Query {query.query_id} not found in result. This should not happen. All query-ids: {result['qid'].unique()}"
+                )
         else:
             temp_result = result[result["qid"] == query.query_id]
             if not temp_result.empty:
-                query.query = temp_result.at[temp_result.index[0], "query"]
+                query.query = temp_result["query"].iloc[0]
             else:
-                warnings.warn(f"Query {query.query_id} not found in result. This should not happen. All query-ids: {result['qid'].unique()}")
+                warnings.warn(
+                    f"Query {query.query_id} not found in result. This should not happen. All query-ids: {result['qid'].unique()}"
+                )
 
         doc_list: List[Document] = []
         for _, entry in result.iterrows():
@@ -238,16 +244,22 @@ def batch_search(
             for query, _ in inputs:
                 temp_result = result[result["qid"] == query.query_id]
                 if not temp_result.empty:
-                    query.query = temp_result.at[temp_result.index[0], t5_rewriter_module.COPY_REWRITTEN_QUERY_COLUMN]
+                    query.query = temp_result[
+                        t5_rewriter_module.COPY_REWRITTEN_QUERY_COLUMN
+                    ].iloc[0]
                 else:
-                    warnings.warn(f"Query {query.query_id} not found in result. This should not happen. All query-ids: {result['qid'].unique()}")
+                    warnings.warn(
+                        f"Query {query.query_id} not found in result. This should not happen. All query-ids: {result['qid'].unique()}"
+                    )
         else:
             for query, _ in inputs:
                 temp_result = result[result["qid"] == query.query_id]
                 if not temp_result.empty:
-                    query.query = temp_result.at[temp_result.index[0], "query"]
+                    query.query = temp_result["query"].iloc[0]
                 else:
-                    warnings.warn(f"Query {query.query_id} not found in result. This should not happen. All query-ids: {result['qid'].unique()}")
+                    warnings.warn(
+                        f"Query {query.query_id} not found in result. This should not happen. All query-ids: {result['qid'].unique()}"
+                    )
 
         contexts: List[Context] = []
         for query, context in inputs:

py_css/models/baseline.py

@@ -7,6 +7,10 @@
 import pyterrier as pt
 from pyterrier_t5 import MonoT5ReRanker, DuoT5ReRanker
 
+import torch
+
+BATCH_SIZE = 128 if torch.cuda.is_available() else 8
+
 
 class Baseline(base_module.Pipeline):
     """
@@ -42,8 +46,8 @@ def __init__(
         t5_qr = t5_rewriter.T5Rewriter()
         bm25 = pt.BatchRetrieve(index, wmodel="BM25", metadata=["docno", "text"])
         self.top_docs = (t5_qr >> bm25, bm25_docs)
-        self.mono_t5 = (MonoT5ReRanker(), mono_t5_docs)
-        self.duo_t5 = (DuoT5ReRanker(), duo_t5_docs)
+        self.mono_t5 = (MonoT5ReRanker(batch_size=BATCH_SIZE), mono_t5_docs)
+        self.duo_t5 = (DuoT5ReRanker(batch_size=BATCH_SIZE), duo_t5_docs)
 
     def transform_input(
         self, query: base_module.Query, context: base_module.Context

py_css/models/baseline_prf.py

@@ -7,6 +7,10 @@
 import pyterrier as pt
 from pyterrier_t5 import MonoT5ReRanker, DuoT5ReRanker
 
+import torch
+
+BATCH_SIZE = 128 if torch.cuda.is_available() else 8
+
 
 class BaselinePRF(base_module.Pipeline):
     """
@@ -50,8 +54,8 @@ def __init__(
         bm25 = pt.BatchRetrieve(index, wmodel="BM25", metadata=["docno", "text"])
         rm3 = pt.rewrite.RM3(index, fb_docs=rm3_fb_docs, fb_terms=rm3_fb_terms)
         self.top_docs = ((bm25 % rm3_fb_docs) >> rm3 >> bm25, bm25_docs)
-        self.mono_t5 = (MonoT5ReRanker(), mono_t5_docs)
-        self.duo_t5 = (DuoT5ReRanker(), duo_t5_docs)
+        self.mono_t5 = (MonoT5ReRanker(batch_size=BATCH_SIZE), mono_t5_docs)
+        self.duo_t5 = (DuoT5ReRanker(batch_size=BATCH_SIZE), duo_t5_docs)
 
     def transform_input(
         self, query: base_module.Query, context: base_module.Context
@@ -78,9 +82,14 @@ def transform(self, query_df: pd.DataFrame) -> pd.DataFrame:
         )
 
         # Now add in the rewritten queries to the top docs
-        top_docs_df = top_docs_df.merge(rewritten_queries_df, on="qid", how="left")
-        # And overwrite the "query" column again
-        top_docs_df["query"] = top_docs_df[t5_rewriter.COPY_REWRITTEN_QUERY_COLUMN]
+        top_docs_df = pt.model.push_queries(top_docs_df, inplace=True)
+        top_docs_df = pd.merge(
+            top_docs_df,
+            rewritten_queries_df[["qid", "rewritten_query"]],
+            on="qid",
+            how="left",
+        )
+        top_docs_df["query"] = top_docs_df["rewritten_query"]
 
         mono_t5_df = self.mono_t5[0].transform(
             top_docs_df.groupby("qid").head(self.mono_t5[1])