Narrowing down the bug

Author: CodingTil

py_css/models/baseline_prf.py

@@ -110,6 +110,11 @@ def transform(self, query_df: pd.DataFrame) -> pd.DataFrame:
         ), f"{unique_qids} != {set(rewritten_queries_df['qid'].unique())}"
 
         top_docs_df = self.top_docs[0].transform(rewritten_queries_df.copy())
+
+        assert unique_qids == set(
+            top_docs_df["qid"].unique()
+        ), f"{unique_qids} != {set(top_docs_df['qid'].unique())}"
+
         top_docs_df = (
             top_docs_df.sort_values(["qid", "score"], ascending=False)
             .groupby("qid")