add models

ku-nlp · Feb 9, 2024 · dc7e00d · dc7e00d
1 parent cda14bd
commit dc7e00d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/llm_judge/common.py b/llm_judge/common.py
@@ -124,10 +124,12 @@ def estimate_cost(self) -> float:
                 enc.encode(self.ref_answer["choices"][0]["turns"][0])
             )
         num_output_tokens = 200  # Estimated from a few samples
-        if self.judge.model == "gpt-4":
+        if self.judge.model in {"gpt-4", "gpt-4-0613"}:
             return (0.03 * num_input_tokens + 0.06 * num_output_tokens) / 1_000
+        elif self.judge.model == "gpt-4-1106-preview":
+            return (0.01 * num_input_tokens + 0.03 * num_output_tokens) / 1_000
         elif self.judge.model == "gpt-3.5-turbo":
-            return (0.001 * num_input_tokens + 0.002 * num_output_tokens) / 1_000
+            return (0.0005 * num_input_tokens + 0.0015 * num_output_tokens) / 1_000
         raise AssertionError
 
     @staticmethod
@@ -212,10 +214,12 @@ def estimate_cost(self) -> float:
                 enc.encode(self.ref_answer["choices"][0]["turns"][0])
             )
         num_output_tokens = 200  # Estimated from a few samples
-        if self.judge.model == "gpt-4":
-            return 2 * (0.03 * num_input_tokens + 0.06 * num_output_tokens) / 1_000
+        if self.judge.model in {"gpt-4", "gpt-4-0613"}:
+            return (0.03 * num_input_tokens + 0.06 * num_output_tokens) / 1_000
+        elif self.judge.model == "gpt-4-1106-preview":
+            return (0.01 * num_input_tokens + 0.03 * num_output_tokens) / 1_000
         elif self.judge.model == "gpt-3.5-turbo":
-            return 2 * (0.001 * num_input_tokens + 0.002 * num_output_tokens) / 1_000
+            return (0.0005 * num_input_tokens + 0.0015 * num_output_tokens) / 1_000
         raise AssertionError
 
     @staticmethod

diff --git a/llm_judge/gen_judgment.py b/llm_judge/gen_judgment.py
@@ -133,7 +133,7 @@ def make_match_groups_pairwise(
         "--judge-model",
         type=str,
         default="gpt-4",
-        choices=["gpt-4", "gpt-3.5-turbo"],
+        choices=["gpt-4", "gpt-4-0613", "gpt-4-1106-preview", "gpt-3.5-turbo"],
         help="The judge model.",
     )
     parser.add_argument(