CAAFE: Add Bedrock Support (#76)

autogluon · Nov 1, 2024 · bdfe008 · bdfe008
1 parent ee0ea2d
commit bdfe008
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 8 deletions.
diff --git a/config/config.yaml b/config/config.yaml
@@ -9,7 +9,10 @@ save_artifacts:
 feature_transformers:
   - _target_: autogluon_assistant.transformer.CAAFETransformer
     eval_model: lightgbm
-    llm_model: gpt-4o-2024-08-06
+    llm_provider: bedrock
+    llm_model: "anthropic.claude-3-5-sonnet-20241022-v2:0"
+    # llm_provider: openai
+    # llm_model: gpt-4o-2024-08-06
     num_iterations: 5
     optimization_metric: roc
   - _target_: autogluon_assistant.transformer.OpenFETransformer

diff --git a/pyproject.toml b/pyproject.toml
@@ -27,7 +27,8 @@ dependencies = [
     "python-calamine",
     "sentence-transformers>=3.1.0",
     "tenacity>=8.2.2,<10.0",
-    "gensim>=4.3", 
+    "gensim>=4.3",
+    "pandas>=2.2"
 ]
 
 [project.optional-dependencies]

diff --git a/src/autogluon_assistant/transformer/feature_transformers/caafe.py b/src/autogluon_assistant/transformer/feature_transformers/caafe.py
@@ -17,45 +17,53 @@ class CAAFETransformer(BaseFeatureTransformer):
 
     def __init__(
         self,
+        llm_provider: str = "openai",
         llm_model: str = "gpt-3.5-turbo",
         num_iterations: int = 2,
         optimization_metric: str = "roc",
         eval_model: str = "lightgbm",
+        region_name: str = "us-west-2",
         **kwargs,
     ) -> None:
-        import openai
+        # Set up credentials if using OpenAI
+        if llm_provider == "openai":
+            import openai
 
-        openai.api_key = kwargs.get("openai_api_key", os.environ.get("OPENAI_API_KEY"))
+            openai.api_key = kwargs.get("openai_api_key", os.environ.get("OPENAI_API_KEY"))
 
         pd.set_option("future.no_silent_downcasting", True)
 
+        self.llm_provider = llm_provider
         self.llm_model = llm_model
         self.iterations = num_iterations
         self.optimization_metric = optimization_metric
         self.eval_model = eval_model
+        self.region_name = region_name
 
+        # Initialize the base classifier
         if self.eval_model == "tab_pfn":
             from tabpfn import TabPFNClassifier
 
             clf_no_feat_eng = TabPFNClassifier(device="cpu", N_ensemble_configurations=16)
-
         elif self.eval_model == "lightgbm":
             from lightgbm import LGBMClassifier
 
             clf_no_feat_eng = LGBMClassifier()
-
         else:
             raise ValueError(f"Unsupported CAAFE eval model: {self.eval_model}")
 
         self.caafe_clf = CAAFEClassifier(
             base_classifier=clf_no_feat_eng,
             optimization_metric=self.optimization_metric,
+            llm_provider=self.llm_provider,
             llm_model=self.llm_model,
             iterations=self.iterations,
+            region_name=self.region_name,
             display_method="print",
+            **kwargs,  # Pass through any additional provider-specific kwargs
         )
 
-        self.metadata = {"transformer": "CAAFE"}
+        self.metadata = {"transformer": "CAAFE", "llm_provider": llm_provider, "llm_model": llm_model}
 
     def _fit_dataframes(
         self,
@@ -83,7 +91,8 @@ def _fit_dataframes(
             train_X.columns,
             target_column_name,
         )
-        logger.info("CAAFE generated features:")
+
+        logger.info(f"CAAFE generated features using {self.llm_provider} model {self.llm_model}:")
         logger.info(self.caafe_clf.code)
 
     def _transform_dataframes(self, train_X: pd.DataFrame, test_X: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]: