config

config.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
 
 from zeno_build.evaluation.text_features.capitalization import input_capital_char_ratio
@@ -25,6 +26,8 @@ from zeno_build.evaluation.text_metrics.critique import (
 )
 from zeno_build.experiments import search_space
 
+from modeling import remove_leading_language
+
 lang_pairs: dict[str, list[str]] = {
     # All language pairs used in any experiment
     "all_lang_pairs": [
@@ -66,13 +69,17 @@ main_space = search_space.CombinatorialSearchSpace(
         "lang_pairs": search_space.Constant("all_lang_pairs"),
         "model_preset": search_space.Categorical(
             [
+                "text-davinci-003-zeroshot",
                 "text-davinci-003-RR-1-shot",
                 "text-davinci-003-RR-5-shot",
                 "text-davinci-003-QR-1-shot",
                 "text-davinci-003-QR-5-shot",
-                "text-davinci-003-zeroshot",
-                "wmt-best",
+                "gpt-3.5-turbo-0301-zeroshot",
+                "gpt-4-0314-zeroshot",
+                "gpt-4-0314-zeroshot-postprocess",
                 "MS-Translator",
+                "google-cloud",
+                "wmt-best",
             ]
         ),
     }
@@ -87,6 +94,7 @@ class GptMtConfig:
     base_model: str
     prompt_strategy: str | None = None
     prompt_shots: int | None = None
+    post_processors: list[Callable[[str], str]] | None = None
 
 
 # The details of each model
@@ -106,8 +114,16 @@ model_configs = {
     "text-davinci-003-zeroshot": GptMtConfig(
         "text-davinci-003/zeroshot", "text-davinci-003", None, 0
     ),
-    "wmt-best": GptMtConfig("wmt-best", "wmt-best"),
+    "gpt-3.5-turbo-0301-zeroshot": GptMtConfig(
+        "gpt-3.5-turbo-0301/zeroshot", "gpt-3.5-turbo-0301", None, 0
+    ),
+    "gpt-4-0314-zeroshot": GptMtConfig("gpt-4-0314/zeroshot", "gpt-4-0314", None, 0),
+    "gpt-4-0314-zeroshot-postprocess": GptMtConfig(
+        "gpt-4-0314/zeroshot", "gpt-4-0314", None, 0, [remove_leading_language]
+    ),
     "MS-Translator": GptMtConfig("MS-Translator", "MS-Translator"),
+    "google-cloud": GptMtConfig("google-cloud", "google-cloud"),
+    "wmt-best": GptMtConfig("wmt-best", "wmt-best"),
 }
 
 sweep_distill_functions = [chrf]