Merge pull request #24 from aws/flex-npo

Feat: Make NovaPromptOptimizer more customizable by changing Prompting Models

Author: VikramShenoy97

src/amzn_nova_prompt_optimizer/__version__.py

@@ -11,4 +11,4 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-VERSION = "1.0.48"  # pragma: no cover
+VERSION = "1.0.52"  # pragma: no cover

src/amzn_nova_prompt_optimizer/core/optimizers/nova_meta_prompter/nova_mp_optimizer.py

@@ -42,7 +42,7 @@ def __init__(self, prompt_adapter: PromptAdapter,
 
 
     def optimize(self, prompter_model_id: str = DEFAULT_PROMPTER_MODEL_ID, max_retries: int = 5):
-        logger.info("Optimizing prompt using Nova Meta Prompter")
+        logger.info(f"Optimizing prompt using Nova Meta Prompter with Model: {prompter_model_id}")
         if not self.inference_adapter:
             raise ValueError("Inference Adapter not passed. "
                              "Initialize and Pass Inference Adapter to use this Optimizer")

src/amzn_nova_prompt_optimizer/core/optimizers/nova_prompt_optimizer/nova_prompt_optimizer.py

@@ -24,29 +24,37 @@
 
 logger = logging.getLogger(__name__)
 
-NOVA_OPTIMA_MODE: Dict[str, Dict[str, Any]] = {
+NOVA_PROMPT_OPTIMIZER_MODE: Dict[str, Dict[str, Any]] = {
     "micro": {
+        "meta_prompt_model_id": "us.amazon.nova-premier-v1:0",
+        "prompter_model_id": "us.amazon.nova-premier-v1:0",
         "task_model_id": "us.amazon.nova-micro-v1:0",
         "num_candidates": 20,
         "num_trials": 30,
         "max_bootstrapped_demos": 4,
         "max_labeled_demos": 4
     },
     "lite": {
+        "meta_prompt_model_id": "us.amazon.nova-premier-v1:0",
+        "prompter_model_id": "us.amazon.nova-premier-v1:0",
         "task_model_id": "us.amazon.nova-lite-v1:0",
         "num_candidates": 20,
         "num_trials": 30,
         "max_bootstrapped_demos": 4,
         "max_labeled_demos": 4
     },
     "pro": {
+        "meta_prompt_model_id": "us.amazon.nova-premier-v1:0",
+        "prompter_model_id": "us.amazon.nova-premier-v1:0",
         "task_model_id": "us.amazon.nova-pro-v1:0",
         "num_candidates": 20,
         "num_trials": 30,
         "max_bootstrapped_demos": 4,
         "max_labeled_demos": 4
     },
     "premier": {
+        "meta_prompt_model_id": "us.amazon.nova-premier-v1:0",
+        "prompter_model_id": "us.amazon.nova-premier-v1:0",
         "task_model_id": "us.amazon.nova-premier-v1:0",
         "num_candidates": 20,
         "num_trials": 30,
@@ -58,10 +66,9 @@
 
 class NovaPromptOptimizer(OptimizationAdapter):
     """
-    NovaOptima is a combination of Meta Prompting and MIPROv2 for Nova Models that yields a stable
+    NovaPromptOptimizer is a combination of Meta Prompting and MIPROv2 for Nova Models that yields a stable
     prompt optimization result.
     """
-
     def __init__(self, prompt_adapter: PromptAdapter,
                  inference_adapter: InferenceAdapter,
                  dataset_adapter: DatasetAdapter,
@@ -75,34 +82,42 @@ def __init__(self, prompt_adapter: PromptAdapter,
         self.meta_prompt_optimization_adapter = NovaMPOptimizationAdapter(prompt_adapter, inference_adapter)
 
     def optimize(self, mode: str = "pro", custom_params = None) -> PromptAdapter:
+        if mode == "custom":
+            if not custom_params:
+                raise ValueError("Custom mode requires custom_params dictionary")
+            required_keys = {"task_model_id", "num_candidates", "num_trials",
+                             "max_bootstrapped_demos", "max_labeled_demos"}
+            if not all(key in custom_params for key in required_keys):
+                raise ValueError(f"custom_params must contain all required keys: {required_keys}")
+            meta_prompt_model_id = custom_params.pop("meta_prompt_model_id", None)
+            optimization_params = custom_params
+        else:
+            if mode not in NOVA_PROMPT_OPTIMIZER_MODE:
+                logger.warning(f"Mode '{mode}' not detected, defaulting to 'pro' mode")
+                mode = "pro"
+            config = NOVA_PROMPT_OPTIMIZER_MODE[mode].copy()  # Create a copy to avoid modifying the original
+            meta_prompt_model_id = config.pop("meta_prompt_model_id")
+            optimization_params = config
+
+
         if not self.inference_adapter:
             raise ValueError("Inference Adapter not passed. "
                              "Initialize and Pass Inference Adapter to use this Optimizer")
-        intermediate_prompt_adapter = self.meta_prompt_optimization_adapter.optimize()
+        if meta_prompt_model_id:
+            intermediate_prompt_adapter = (
+                self.meta_prompt_optimization_adapter.optimize(prompter_model_id=meta_prompt_model_id))
+        else:
+            intermediate_prompt_adapter = self.meta_prompt_optimization_adapter.optimize()
+
         if not self.dataset_adapter or not self.metric_adapter:
-            logger.info("[Nova Optima] No Dataset or No metric provided, running only Nova Meta Prompter")
+            logger.info("[Nova Prompt Optimizer] No Dataset or No metric provided, running only Nova Meta Prompter")
             return intermediate_prompt_adapter
 
-        nova_optima_optimization_adapter = NovaMIPROv2OptimizationAdapter(
+        nova_prompt_optimizer = NovaMIPROv2OptimizationAdapter(
             prompt_adapter=intermediate_prompt_adapter,
             dataset_adapter=self.dataset_adapter,
             metric_adapter=self.metric_adapter,
             inference_adapter=self.inference_adapter)
 
-        if mode == "custom":
-            if not custom_params:
-                raise ValueError("Custom mode requires custom_params dictionary")
-            required_keys = {"task_model_id", "num_candidates", "num_trials",
-                             "max_bootstrapped_demos", "max_labeled_demos"}
-            if not all(key in custom_params for key in required_keys):
-                raise ValueError(f"custom_params must contain all required keys: {required_keys}")
-            optimization_params = custom_params
-        else:
-            if mode not in NOVA_OPTIMA_MODE:
-                logger.warning(f"Mode '{mode}' not detected, defaulting to 'pro' mode")
-                optimization_params = NOVA_OPTIMA_MODE["pro"]
-            else:
-                optimization_params = NOVA_OPTIMA_MODE[mode]
-        optimized_prompt_adapter = nova_optima_optimization_adapter.optimize(**optimization_params,
-                                                                             enable_json_fallback=False)
+        optimized_prompt_adapter = nova_prompt_optimizer.optimize(**optimization_params, enable_json_fallback=False)
         return optimized_prompt_adapter