Make settings explicit

.github/tests/lm_tests.py

@@ -288,12 +288,7 @@ def test_filter_cascade(setup_models):
 def test_join_cascade(setup_models):
     models = setup_models
     rm = SentenceTransformersRM(model="intfloat/e5-base-v2")
-    lotus.settings.configure(
-        lm=models["gpt-4o-mini"],
-        rm=rm,
-        min_join_cascade_size=10,  # for smaller testings
-        cascade_IS_random_seed=42,
-    )
+    lotus.settings.configure(lm=models["gpt-4o-mini"], rm=rm)
 
     data1 = {
         "School": [
@@ -326,7 +321,12 @@ def test_join_cascade(setup_models):
 
     # Cascade join
     joined_df, stats = df1.sem_join(
-        df2, join_instruction, cascade_args=CascadeArgs(recall_target=0.7, precision_target=0.7), return_stats=True
+        df2,
+        join_instruction,
+        cascade_args=CascadeArgs(
+            recall_target=0.7, precision_target=0.7, min_join_cascade_size=10, cascade_IS_random_seed=42
+        ),
+        return_stats=True,
     )
 
     for pair in expected_pairs:
@@ -337,7 +337,12 @@ def test_join_cascade(setup_models):
 
     # All joins resolved by the large model
     joined_df, stats = df1.sem_join(
-        df2, join_instruction, cascade_args=CascadeArgs(recall_target=1.0, precision_target=1.0), return_stats=True
+        df2,
+        join_instruction,
+        cascade_args=CascadeArgs(
+            recall_target=1.0, precision_target=1.0, min_join_cascade_size=10, cascade_IS_random_seed=42
+        ),
+        return_stats=True,
     )
 
     for pair in expected_pairs:

examples/op_examples/filter_cascade.py

@@ -4,11 +4,10 @@
 from lotus.models import LM
 from lotus.types import CascadeArgs
 
-
-gpt_35_turbo = LM("gpt-3.5-turbo")
+gpt_4o_mini = LM("gpt-4o-mini")
 gpt_4o = LM("gpt-4o")
 
-lotus.settings.configure(lm=gpt_4o, helper_lm=gpt_35_turbo)
+lotus.settings.configure(lm=gpt_4o, helper_lm=gpt_4o_mini)
 data = {
     "Course Name": [
         "Probability and Random Processes",

lotus/sem_ops/cascade_utils.py

@@ -2,42 +2,39 @@
 from numpy.typing import NDArray
 
 import lotus
+from lotus.types import CascadeArgs
 
 
 def importance_sampling(
     proxy_scores: list[float],
-    sample_percentage: float,
+    cascade_args: CascadeArgs,
 ) -> tuple[NDArray[np.int64], NDArray[np.float64]]:
     """Uses importance sampling and returns the list of indices from which to learn cascade thresholds."""
-    if lotus.settings.cascade_IS_random_seed is not None:
-        np.random.seed(lotus.settings.cascade_IS_random_seed)
+    if cascade_args.cascade_IS_random_seed is not None:
+        np.random.seed(cascade_args.cascade_IS_random_seed)
 
     w = np.sqrt(proxy_scores)
-    is_weight = lotus.settings.cascade_IS_weight
+    is_weight = cascade_args.cascade_IS_weight
     w = is_weight * w / np.sum(w) + (1 - is_weight) * np.ones((len(proxy_scores))) / len(proxy_scores)
 
-    if lotus.settings.cascade_IS_max_sample_range is not None:
-        sample_range = min(lotus.settings.cascade_IS_max_sample_range, len(proxy_scores))
-        sample_w = w[:sample_range]
-        sample_w = sample_w / np.sum(sample_w)
-        indices = np.arange(sample_range)
-    else:
-        sample_w = w
-        indices = np.arange(len(proxy_scores))
+    sample_range = min(cascade_args.cascade_IS_max_sample_range, len(proxy_scores))
+    sample_w = w[:sample_range]
+    sample_w = sample_w / np.sum(sample_w)
+    indices = np.arange(sample_range)
 
-    sample_size = int(sample_percentage * len(proxy_scores))
+    sample_size = int(cascade_args.sampling_percentage * len(proxy_scores))
     sample_indices = np.random.choice(indices, sample_size, p=sample_w)
 
     correction_factors = (1 / len(proxy_scores)) / w
 
     return sample_indices, correction_factors
 
 
-def calibrate_llm_logprobs(true_probs: list[float]) -> list[float]:
+def calibrate_llm_logprobs(true_probs: list[float], cascade_args: CascadeArgs) -> list[float]:
     """Transforms true probabilities to calibrate LLM proxies."""
-    num_quantiles = lotus.settings.cascade_num_calibration_quantiles
+    num_quantiles = cascade_args.cascade_num_calibration_quantiles
     quantile_values = np.percentile(true_probs, np.linspace(0, 100, num_quantiles + 1))
-    true_probs = (np.digitize(true_probs, quantile_values) - 1) / num_quantiles
+    true_probs = list((np.digitize(true_probs, quantile_values) - 1) / num_quantiles)
     true_probs = list(np.clip(true_probs, 0, 1))
     return true_probs
 
@@ -46,9 +43,7 @@ def learn_cascade_thresholds(
     proxy_scores: list[float],
     oracle_outputs: list[bool],
     sample_correction_factors: NDArray[np.float64],
-    recall_target: float,
-    precision_target: float,
-    delta: float,
+    cascade_args: CascadeArgs,
 ) -> tuple[tuple[float, float], int]:
     """Learns cascade thresholds given targets and proxy scores,
     oracle outputs over the sample, and correction factors for the
@@ -97,7 +92,7 @@ def calculate_tau_neg(sorted_pairs: list[tuple[float, bool, float]], tau_pos: fl
     best_combination = (1.0, 0.0)  # initial tau_+, tau_-
 
     # Find tau_negative based on recall
-    tau_neg_0 = calculate_tau_neg(sorted_pairs, best_combination[0], recall_target)
+    tau_neg_0 = calculate_tau_neg(sorted_pairs, best_combination[0], cascade_args.recall_target)
     best_combination = (best_combination[0], tau_neg_0)
 
     # Do a statistical correction to get a new target recall
@@ -109,8 +104,8 @@ def calculate_tau_neg(sorted_pairs: list[tuple[float, bool, float]], tau_pos: fl
     mean_z2 = float(np.mean(Z2)) if Z2 else 0.0
     std_z2 = float(np.std(Z2)) if Z2 else 0.0
 
-    ub_z1 = UB(mean_z1, std_z1, sample_size, delta / 2)
-    lb_z2 = LB(mean_z2, std_z2, sample_size, delta / 2)
+    ub_z1 = UB(mean_z1, std_z1, sample_size, cascade_args.failure_probability / 2)
+    lb_z2 = LB(mean_z2, std_z2, sample_size, cascade_args.failure_probability / 2)
     if ub_z1 + lb_z2 == 0:  # Avoid division by zero
         corrected_recall_target = 1.0
     else:
@@ -127,8 +122,8 @@ def calculate_tau_neg(sorted_pairs: list[tuple[float, bool, float]], tau_pos: fl
         Z = [int(x[1]) for x in sorted_pairs if x[0] >= possible_threshold]
         mean_z = float(np.mean(Z)) if Z else 0.0
         std_z = float(np.std(Z)) if Z else 0.0
-        p_l = LB(mean_z, std_z, len(Z), delta / len(sorted_pairs))
-        if p_l > precision_target:
+        p_l = LB(mean_z, std_z, len(Z), cascade_args.failure_probability / len(sorted_pairs))
+        if p_l > cascade_args.precision_target:
             candidate_thresholds.append(possible_threshold)
 
     best_combination = (max(best_combination[1], min(candidate_thresholds)), best_combination[1])

lotus/sem_ops/sem_agg.py

@@ -164,6 +164,11 @@ def __call__(
             pd.DataFrame: The dataframe with the aggregated answer.
         """
 
+        if lotus.settings.lm is None:
+            raise ValueError(
+                "The language model must be an instance of LM. Please configure a valid language model using lotus.settings.configure()"
+            )
+
         lotus.logger.debug(f"User instruction: {user_instruction}")
         if all_cols:
             col_li = list(self._obj.columns)

lotus/sem_ops/sem_cluster_by.py

@@ -37,6 +37,11 @@ def __call__(
         Returns:
             pd.DataFrame: The DataFrame with the cluster assignments.
         """
+        if lotus.settings.rm is None:
+            raise ValueError(
+                "The retrieval model must be an instance of RM. Please configure a valid retrieval model using lotus.settings.configure()"
+            )
+
         cluster_fn = lotus.utils.cluster(col_name, ncentroids)
         indices = cluster_fn(self._obj, niter, verbose)
 

lotus/sem_ops/sem_dedup.py

@@ -34,6 +34,11 @@ def __call__(
         Returns:
             pd.DataFrame: The DataFrame with duplicates removed.
         """
+        if lotus.settings.rm is None:
+            raise ValueError(
+                "The retrieval model must be an instance of RM. Please configure a valid retrieval model using lotus.settings.configure()"
+            )
+
         joined_df = self._obj.sem_sim_join(self._obj, col_name, col_name, len(self._obj), lsuffix="_l", rsuffix="_r")
         dedup_df = joined_df[joined_df["_scores"] > threshold]
         dedup_df = dedup_df[dedup_df[f"{col_name}_l"] != dedup_df[f"{col_name}_r"]]

lotus/sem_ops/sem_extract.py

@@ -95,6 +95,10 @@ def __call__(
         Returns:
             pd.DataFrame: The dataframe with the new mapped columns.
         """
+        if lotus.settings.lm is None:
+            raise ValueError(
+                "The language model must be an instance of LM. Please configure a valid language model using lotus.settings.configure()"
+            )
 
         # check that column exists
         for column in input_cols:

lotus/sem_ops/sem_filter.py

@@ -79,9 +79,7 @@ def learn_filter_cascade_thresholds(
     lm: lotus.models.LM,
     formatted_usr_instr: str,
     default: bool,
-    recall_target: float,
-    precision_target: float,
-    delta: float,
+    cascade_args: CascadeArgs,
     helper_true_probs: list[float],
     sample_correction_factors: NDArray[np.float64],
     examples_multimodal_data: list[dict[str, Any]] | None = None,
@@ -111,9 +109,7 @@ def learn_filter_cascade_thresholds(
             proxy_scores=helper_true_probs,
             oracle_outputs=large_outputs,
             sample_correction_factors=sample_correction_factors,
-            recall_target=recall_target,
-            precision_target=precision_target,
-            delta=delta,
+            cascade_args=cascade_args,
         )
 
         lotus.logger.info(f"Learned cascade thresholds: {best_combination}")
@@ -174,6 +170,11 @@ def __call__(
         Returns:
             pd.DataFrame | tuple[pd.DataFrame, dict[str, Any]]: The filtered dataframe or a tuple containing the filtered dataframe and statistics.
         """
+        if lotus.settings.lm is None:
+            raise ValueError(
+                "The language model must be an instance of LM. Please configure a valid language model using lotus.settings.configure()"
+            )
+
         stats = {}
         lotus.logger.debug(user_instruction)
         col_li = lotus.nl_expression.parse_cols(user_instruction)
@@ -246,14 +247,13 @@ def __call__(
                 progress_bar_desc="Running helper LM",
             )
             helper_outputs, helper_logprobs = helper_output.outputs, helper_output.logprobs
+            assert helper_logprobs is not None
             formatted_helper_logprobs: LogprobsForFilterCascade = (
                 lotus.settings.helper_lm.format_logprobs_for_filter_cascade(helper_logprobs)
             )
-            helper_true_probs = calibrate_llm_logprobs(formatted_helper_logprobs.true_probs)
+            helper_true_probs = calibrate_llm_logprobs(formatted_helper_logprobs.true_probs, cascade_args)
 
-            sample_indices, correction_factors = importance_sampling(
-                helper_true_probs, cascade_args.sampling_percentage
-            )
+            sample_indices, correction_factors = importance_sampling(helper_true_probs, cascade_args)
             sample_df = self._obj.loc[sample_indices]
             sample_multimodal_data = task_instructions.df2multimodal_info(sample_df, col_li)
             sample_helper_true_probs = [helper_true_probs[i] for i in sample_indices]
@@ -264,9 +264,7 @@ def __call__(
                 lm=lotus.settings.lm,
                 formatted_usr_instr=formatted_usr_instr,
                 default=default,
-                recall_target=cascade_args.recall_target,
-                precision_target=cascade_args.precision_target,
-                delta=cascade_args.failure_probability / 2,
+                cascade_args=cascade_args,
                 helper_true_probs=sample_helper_true_probs,
                 sample_correction_factors=sample_correction_factors,
                 examples_multimodal_data=examples_multimodal_data,

lotus/sem_ops/sem_index.py

@@ -30,9 +30,12 @@ def __call__(self, col_name: str, index_dir: str) -> pd.DataFrame:
         Returns:
             pd.DataFrame: The DataFrame with the index directory saved.
         """
+        if lotus.settings.rm is None:
+            raise ValueError(
+                "The retrieval model must be an instance of RM. Please configure a valid retrieval model using lotus.settings.configure()"
+            )
+
         rm = lotus.settings.rm
-        if rm is None:
-            raise AttributeError("Must set rm in lotus.settings")
         rm.index(self._obj[col_name], index_dir)
         self._obj.attrs["index_dirs"][col_name] = index_dir
         return self._obj