Minor Fix

Author: CodingTil

py_css/models/base.py

@@ -2,6 +2,7 @@
 from dataclasses import dataclass
 import logging
 from typing import Optional, List, Tuple, TypeAlias
+import warnings
 
 import models.T5Rewriter as t5_rewriter_module
 
@@ -189,11 +190,17 @@ def search(self, query: Query, context: Context) -> Tuple[Context, pd.DataFrame]
         result = self.transform(query_df)
 
         if t5_rewriter_module.COPY_REWRITTEN_QUERY_COLUMN in result.columns:
-            query.query = result[t5_rewriter_module.COPY_REWRITTEN_QUERY_COLUMN].iloc[0]
-            print("QUERY COMES FROM REWRITER")
+            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]
+            else:
+                warnings.warn(f"Query {query.query_id} not found in result. This should not happen. All query-ids: {result['qid'].unique()}")
         else:
-            query.query = result["query"].iloc[0]
-            print("QUERY COMES FROM ORIGINAL?!?!?!?!")
+            temp_result = result[result["qid"] == query.query_id]
+            if not temp_result.empty:
+                query.query = temp_result.at[temp_result.index[0], "query"]
+            else:
+                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():
@@ -229,12 +236,18 @@ def batch_search(
 
         if t5_rewriter_module.COPY_REWRITTEN_QUERY_COLUMN in result.columns:
             for query, _ in inputs:
-                query.query = result[result["qid"] == query.query_id][
-                    t5_rewriter_module.COPY_REWRITTEN_QUERY_COLUMN
-                ].iloc[0]
+                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]
+                else:
+                    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:
-                query.query = result[result["qid"] == query.query_id]["query"].iloc[0]
+                temp_result = result[result["qid"] == query.query_id]
+                if not temp_result.empty:
+                    query.query = temp_result.at[temp_result.index[0], "query"]
+                else:
+                    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: