set orig_index in Index.__call__ to avoid pandas SettingWithCopyWarning

mrjleo · Oct 18, 2024 · 7c94a98 · 7c94a98
1 parent 4308869
commit 7c94a98
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/fast_forward/index/__init__.py b/fast_forward/index/__init__.py
@@ -369,9 +369,6 @@ def _early_stopping(
         Returns:
             pd.DataFrame: Data frame with computed scores.
         """
-        # early stopping splits the data frame, hence we need to keep track of the original index
-        df["orig_index"] = df.index
-
         # data frame for computed scores
         scores_so_far = None
 
@@ -464,6 +461,9 @@ def __call__(
         query_df["q_no"] = query_df.index
         df_with_q_no = ranking._df.merge(query_df, on="q_id", suffixes=[None, "_"])
 
+        # early stopping splits the data frame, hence we need to keep track of the original index
+        df_with_q_no["orig_index"] = df_with_q_no.index
+
         # batch encode queries
         query_vectors = self.encode_queries(list(query_df["query"]))