diff --git a/scripts/estimate_cost.py b/scripts/estimate_cost.py index 5d7797df86..8442c0c197 100644 --- a/scripts/estimate_cost.py +++ b/scripts/estimate_cost.py @@ -17,17 +17,22 @@ class ModelCost: total_num_prompt_tokens: int = 0 - total_max_num_completion_tokens: int = 0 + total_num_completion_tokens: int = 0 + + total_num_instances: int = 0 @property def total_tokens(self) -> int: - return self.total_num_prompt_tokens + self.total_max_num_completion_tokens + return self.total_num_prompt_tokens + self.total_num_completion_tokens def add_prompt_tokens(self, num_tokens: int): self.total_num_prompt_tokens += num_tokens def add_num_completion_tokens(self, num_tokens: int): - self.total_max_num_completion_tokens += num_tokens + self.total_num_completion_tokens += num_tokens + + def add_num_instances(self, num_instances: int): + self.total_num_instances += num_instances class CostCalculator: @@ -60,10 +65,17 @@ def aggregate(self) -> Dict[str, ModelCost]: for metric in metrics: cost: ModelCost = models_to_costs[model] metric_name: str = metric["name"]["name"] + + # Don't count perturbations + if "perturbation" in metric["name"]: + continue + if metric_name == "num_prompt_tokens": cost.add_prompt_tokens(metric["sum"]) - elif metric_name == "max_num_completion_tokens": + elif metric_name == "num_completion_tokens": cost.add_num_completion_tokens(metric["sum"]) + elif metric_name == "num_instances": + cost.add_num_instances(metric["sum"]) return models_to_costs @@ -78,6 +90,6 @@ def aggregate(self) -> Dict[str, ModelCost]: for model_name, model_cost in model_costs.items(): print( f"{model_name}: Total prompt tokens={model_cost.total_num_prompt_tokens} + " - f"Total max completion tokens={model_cost.total_max_num_completion_tokens} = " - f"{model_cost.total_tokens}" + f"Total completion tokens={model_cost.total_num_completion_tokens} = " + f"{model_cost.total_tokens} for {model_cost.total_num_instances} instances." )