operator level cache

lotus/cache.py

@@ -1,3 +1,5 @@
+import hashlib
+import json
 import os
 import pickle
 import sqlite3
@@ -8,6 +10,8 @@
 from functools import wraps
 from typing import Any, Callable
 
+import pandas as pd
+
 import lotus
 
 
@@ -23,6 +27,44 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
+def operator_cache(func: Callable) -> Callable:
+    """Decorator to add operator level caching."""
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        model = lotus.settings.lm
+        use_operator_cache = kwargs.get("use_operator_cache", False)
+
+        if use_operator_cache and model.cache:
+
+            def serialize(value):
+                if isinstance(value, pd.DataFrame):
+                    return value.to_json()
+                elif hasattr(value, "dict"):
+                    return value.dict()
+                return value
+
+            serialized_kwargs = {key: serialize(value) for key, value in kwargs.items()}
+            serialized_args = [serialize(arg) for arg in args]
+            cache_key = hashlib.sha256(
+                json.dumps({"args": serialized_args, "kwargs": serialized_kwargs}, sort_keys=True).encode()
+            ).hexdigest()
+
+            cached_result = model.cache.get(cache_key)
+            if cached_result is not None:
+                print(f"Cache hit for {cache_key}")
+                return cached_result
+            print(f"Cache miss for {cache_key}")
+
+            result = func(self, *args, **kwargs)
+            model.cache.insert(cache_key, result)
+            return result
+
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
+
 class CacheType(Enum):
     IN_MEMORY = "in_memory"
     SQLITE = "sqlite"

lotus/sem_ops/sem_agg.py

@@ -3,6 +3,7 @@
 import pandas as pd
 
 import lotus.models
+from lotus.cache import operator_cache
 from lotus.templates import task_instructions
 from lotus.types import LMOutput, SemanticAggOutput
 
@@ -14,6 +15,7 @@ def sem_agg(
     partition_ids: list[int],
     safe_mode: bool = False,
     progress_bar_desc: str = "Aggregating",
+    use_operator_cache: bool = False,
 ) -> SemanticAggOutput:
     """
     Aggregates multiple documents into a single answer using a model.
@@ -148,6 +150,7 @@ def process_group(args):
         group, user_instruction, all_cols, suffix, progress_bar_desc = args
         return group.sem_agg(user_instruction, all_cols, suffix, None, progress_bar_desc=progress_bar_desc)
 
+    @operator_cache
     def __call__(
         self,
         user_instruction: str,
@@ -156,6 +159,7 @@ def __call__(
         group_by: list[str] | None = None,
         safe_mode: bool = False,
         progress_bar_desc: str = "Aggregating",
+        use_operator_cache: bool = False,
     ) -> pd.DataFrame:
         """
         Applies semantic aggregation over a dataframe.
@@ -213,6 +217,7 @@ def __call__(
             partition_ids,
             safe_mode=safe_mode,
             progress_bar_desc=progress_bar_desc,
+            use_operator_cache=use_operator_cache,
         )
 
         # package answer in a dataframe

lotus/sem_ops/sem_cluster_by.py

@@ -4,6 +4,7 @@
 import pandas as pd
 
 import lotus
+from lotus.cache import operator_cache
 
 
 @pd.api.extensions.register_dataframe_accessor("sem_cluster_by")
@@ -19,6 +20,7 @@ def _validate(obj: Any) -> None:
         if not isinstance(obj, pd.DataFrame):
             raise AttributeError("Must be a DataFrame")
 
+    @operator_cache
     def __call__(
         self,
         col_name: str,
@@ -27,6 +29,7 @@ def __call__(
         return_centroids: bool = False,
         niter: int = 20,
         verbose: bool = False,
+        use_operator_cache: bool = False,
     ) -> pd.DataFrame | tuple[pd.DataFrame, np.ndarray]:
         """
         Perform semantic clustering on the DataFrame.
@@ -52,7 +55,7 @@ def __call__(
         self._obj["cluster_id"] = pd.Series(indices, index=self._obj.index)
         # if return_scores:
         #     self._obj["centroid_sim_score"] = pd.Series(scores, index=self._obj.index)
-        
+
         # if return_centroids:
         #     return self._obj, centroids
         # else:

lotus/sem_ops/sem_extract.py

@@ -3,6 +3,7 @@
 import pandas as pd
 
 import lotus
+from lotus.cache import operator_cache
 from lotus.models import LM
 from lotus.templates import task_instructions
 from lotus.types import LMOutput, SemanticExtractOutput, SemanticExtractPostprocessOutput
@@ -19,6 +20,7 @@ def sem_extract(
     postprocessor: Callable[[list[str]], SemanticExtractPostprocessOutput] = extract_postprocess,
     safe_mode: bool = False,
     progress_bar_desc: str = "Extracting",
+    use_operator_cache: bool = False,
 ) -> SemanticExtractOutput:
     """
     Extracts attributes and values from a list of documents using a model.
@@ -33,7 +35,6 @@ def sem_extract(
     Returns:
         SemanticExtractOutput: The outputs, raw outputs, and quotes.
     """
-
     # prepare model inputs
     inputs = []
     for doc in docs:
@@ -72,6 +73,7 @@ def _validate(obj: pd.DataFrame) -> None:
         if not isinstance(obj, pd.DataFrame):
             raise AttributeError("Must be a DataFrame")
 
+    @operator_cache
     def __call__(
         self,
         input_cols: list[str],
@@ -81,6 +83,7 @@ def __call__(
         return_raw_outputs: bool = False,
         safe_mode: bool = False,
         progress_bar_desc: str = "Extracting",
+        use_operator_cache: bool = False,
     ) -> pd.DataFrame:
         """
         Extracts the attributes and values of a dataframe.
@@ -115,6 +118,7 @@ def __call__(
             postprocessor=postprocessor,
             safe_mode=safe_mode,
             progress_bar_desc=progress_bar_desc,
+            use_operator_cache=use_operator_cache,
         )
 
         new_df = self._obj.copy()

lotus/sem_ops/sem_filter.py

@@ -5,6 +5,7 @@
 from numpy.typing import NDArray
 
 import lotus
+from lotus.cache import operator_cache
 from lotus.templates import task_instructions
 from lotus.types import CascadeArgs, LMOutput, LogprobsForFilterCascade, SemanticFilterOutput
 from lotus.utils import show_safe_mode
@@ -26,6 +27,7 @@ def sem_filter(
     safe_mode: bool = False,
     show_progress_bar: bool = True,
     progress_bar_desc: str = "Filtering",
+    use_operator_cache: bool = False,
 ) -> SemanticFilterOutput:
     """
     Filters a list of documents based on a given user instruction using a language model.
@@ -103,6 +105,7 @@ def learn_filter_cascade_thresholds(
             strategy=strategy,
             safe_mode=False,
             progress_bar_desc="Running oracle for threshold learning",
+            use_operator_cache=False,
         ).outputs
 
         best_combination, _ = learn_cascade_thresholds(
@@ -134,6 +137,7 @@ def _validate(obj: Any) -> None:
         if not isinstance(obj, pd.DataFrame):
             raise AttributeError("Must be a DataFrame")
 
+    @operator_cache
     def __call__(
         self,
         user_instruction: str,
@@ -148,6 +152,7 @@ def __call__(
         return_stats: bool = False,
         safe_mode: bool = False,
         progress_bar_desc: str = "Filtering",
+        use_operator_cache: bool = False,
     ) -> pd.DataFrame | tuple[pd.DataFrame, dict[str, Any]]:
         """
         Applies semantic filter over a dataframe.
@@ -245,6 +250,7 @@ def __call__(
                 safe_mode=safe_mode,
                 show_progress_bar=True,
                 progress_bar_desc="Running helper LM",
+                use_operator_cache=use_operator_cache,
             )
             helper_outputs, helper_logprobs = helper_output.outputs, helper_output.logprobs
             assert helper_logprobs is not None
@@ -325,6 +331,7 @@ def __call__(
                     strategy=strategy,
                     safe_mode=safe_mode,
                     progress_bar_desc="Running predicate evals with oracle LM",
+                    use_operator_cache=use_operator_cache,
                 )
 
                 for idx, large_idx in enumerate(low_conf_idxs):
@@ -348,6 +355,7 @@ def __call__(
                 safe_mode=safe_mode,
                 show_progress_bar=True,
                 progress_bar_desc=progress_bar_desc,
+                use_operator_cache=use_operator_cache,
             )
             outputs = output.outputs
             raw_outputs = output.raw_outputs

lotus/sem_ops/sem_join.py

@@ -4,6 +4,7 @@
 from tqdm import tqdm
 
 import lotus
+from lotus.cache import operator_cache
 from lotus.templates import task_instructions
 from lotus.types import CascadeArgs, SemanticJoinOutput
 from lotus.utils import show_safe_mode
@@ -29,6 +30,7 @@ def sem_join(
     safe_mode: bool = False,
     show_progress_bar: bool = True,
     progress_bar_desc: str = "Join comparisons",
+    use_operator_cache: bool = False,
 ) -> SemanticJoinOutput:
     """
     Joins two series using a model.
@@ -90,6 +92,7 @@ def sem_join(
             default=default,
             strategy=strategy,
             show_progress_bar=False,
+            use_operator_cache=use_operator_cache,
         )
         outputs = output.outputs
         raw_outputs = output.raw_outputs
@@ -139,6 +142,7 @@ def sem_join_cascade(
     default: bool = True,
     strategy: str | None = None,
     safe_mode: bool = False,
+    use_operator_cache: bool = False,
 ) -> SemanticJoinOutput:
     """
     Joins two series using a cascade helper model and a oracle model.
@@ -235,6 +239,7 @@ def sem_join_cascade(
             default=default,
             strategy=strategy,
             show_progress_bar=False,
+            use_operator_cache=use_operator_cache,
         )
         pbar.update(num_large)
         pbar.close()
@@ -513,6 +518,7 @@ def learn_join_cascade_threshold(
             cot_reasoning=cot_reasoning,
             strategy=strategy,
             progress_bar_desc="Running oracle for threshold learning",
+            use_operator_cache=False,
         )
 
         (pos_threshold, neg_threshold), _ = learn_cascade_thresholds(
@@ -545,6 +551,7 @@ def _validate(obj: Any) -> None:
         if not isinstance(obj, pd.DataFrame):
             raise AttributeError("Must be a DataFrame")
 
+    @operator_cache
     def __call__(
         self,
         other: pd.DataFrame | pd.Series,
@@ -559,6 +566,7 @@ def __call__(
         return_stats: bool = False,
         safe_mode: bool = False,
         progress_bar_desc: str = "Join comparisons",
+        use_operator_cache: bool = False,
     ) -> pd.DataFrame:
         """
         Applies semantic join over a dataframe.
@@ -672,6 +680,7 @@ def __call__(
                 default=default,
                 strategy=strategy,
                 safe_mode=safe_mode,
+                use_operator_cache=use_operator_cache,
             )
         else:
             output = sem_join(
@@ -690,6 +699,7 @@ def __call__(
                 strategy=strategy,
                 safe_mode=safe_mode,
                 progress_bar_desc=progress_bar_desc,
+                use_operator_cache=use_operator_cache,
             )
         join_results = output.join_results
         all_raw_outputs = output.all_raw_outputs

lotus/sem_ops/sem_map.py

@@ -3,6 +3,7 @@
 import pandas as pd
 
 import lotus
+from lotus.cache import operator_cache
 from lotus.templates import task_instructions
 from lotus.types import LMOutput, SemanticMapOutput, SemanticMapPostprocessOutput
 from lotus.utils import show_safe_mode
@@ -21,6 +22,7 @@ def sem_map(
     strategy: str | None = None,
     safe_mode: bool = False,
     progress_bar_desc: str = "Mapping",
+    use_operator_cache: bool = False,
 ) -> SemanticMapOutput:
     """
     Maps a list of documents to a list of outputs using a model.
@@ -80,6 +82,7 @@ def _validate(obj: pd.DataFrame) -> None:
         if not isinstance(obj, pd.DataFrame):
             raise AttributeError("Must be a DataFrame")
 
+    @operator_cache
     def __call__(
         self,
         user_instruction: str,
@@ -91,6 +94,7 @@ def __call__(
         strategy: str | None = None,
         safe_mode: bool = False,
         progress_bar_desc: str = "Mapping",
+        use_operator_cache: bool = False,
     ) -> pd.DataFrame:
         """
         Applies semantic map over a dataframe.
@@ -145,6 +149,7 @@ def __call__(
             strategy=strategy,
             safe_mode=safe_mode,
             progress_bar_desc=progress_bar_desc,
+            use_operator_cache=use_operator_cache,
         )
 
         new_df = self._obj.copy()

lotus/sem_ops/sem_search.py

@@ -3,6 +3,7 @@
 import pandas as pd
 
 import lotus
+from lotus.cache import operator_cache
 from lotus.types import RerankerOutput, RMOutput
 
 
@@ -19,6 +20,7 @@ def _validate(obj: Any) -> None:
         if not isinstance(obj, pd.DataFrame):
             raise AttributeError("Must be a DataFrame")
 
+    @operator_cache
     def __call__(
         self,
         col_name: str,
@@ -27,6 +29,7 @@ def __call__(
         n_rerank: int | None = None,
         return_scores: bool = False,
         suffix: str = "_sim_score",
+        use_operator_cache: bool = False,
     ) -> pd.DataFrame:
         """
         Perform semantic search on the DataFrame.