[WIP] in-loop task evaluator compatible from oe-eval-internal

pyproject.toml

@@ -15,6 +15,7 @@ license = { file = "LICENSE" }
 dependencies = [
     "numpy<2.0",
     "torch>=2.4,<=2.6.0.dev20241009",
+    "torchmetrics<=1.5.1",
     "cached-path",
     "requests",
     "packaging",

src/eval/oe_eval_tasks/pubmedqa/mc_3shot/config.json

@@ -0,0 +1 @@
+{"task_name": "pubmedqa:mc", "task_hash": "6f91fd6326a60b11e8c2d33d632467c0", "task_config": {"task_name": "pubmedqa:mc", "task_core": "pubmedqa", "limit": 1000, "split": "test", "num_shots": 3, "fewshot_seed": 1234, "primary_metric": "acc_raw", "random_subsample_seed": 1234, "context_kwargs": {}, "generation_kwargs": {}, "metric_kwargs": {}, "native_id_field": "id", "fewshot_source": "SciLLM:PubMedQA", "dataset_path": "bigbio/pubmed_qa", "dataset_name": "pubmed_qa_labeled_fold0_source", "use_chat_format": null, "version": 0, "revision": null, "metadata": {"description": "PubMedQA (MC) using SciLLM", "regimes": ["SciLLM-v0.1"], "alias": "pubmedqa:mc::scillm"}}, "current_date": "2024-10-13 03:45:58 UTC", "num_instances": 500}

src/eval/oe_eval_tasks/pubmedqa/rc_3shot/config.json

@@ -0,0 +1 @@
+{"task_name": "pubmedqa", "task_hash": "7145b0916b8d74ab75e87e2768c233df", "task_config": {"task_name": "pubmedqa", "task_core": "pubmedqa", "limit": 1000, "split": "test", "num_shots": 3, "fewshot_seed": 1234, "primary_metric": "acc_raw", "random_subsample_seed": 1234, "context_kwargs": {"description": null}, "generation_kwargs": {}, "metric_kwargs": {"uncond_docid_offset": 1000000}, "native_id_field": "id", "fewshot_source": "SciLLM:PubMedQA", "dataset_path": "bigbio/pubmed_qa", "dataset_name": "pubmed_qa_labeled_fold0_source", "use_chat_format": null, "version": 0, "revision": null, "metadata": {"description": "PubMedQA (RC) using SciLLM", "regimes": ["SciLLM-v0.1"], "alias": "pubmedqa:rc::scillm"}}, "current_date": "2024-10-13 03:45:51 UTC", "num_instances": 500}

src/eval/oe_eval_tasks/scifact/mc_3shot/config.json

@@ -0,0 +1 @@
+{"task_name": "scifact:mc", "task_hash": "5daebde10a80d12c8c6df755e1d5603f", "task_config": {"task_name": "scifact:mc", "task_core": "scifact", "limit": 1000, "split": "validation", "num_shots": 3, "fewshot_seed": 1234, "primary_metric": "acc_raw", "random_subsample_seed": 1234, "context_kwargs": {"description": null}, "generation_kwargs": {}, "metric_kwargs": {}, "native_id_field": "id", "fewshot_source": "SciLLM:SciFact", "dataset_path": "allenai/scifact_entailment", "dataset_name": null, "use_chat_format": null, "version": 0, "revision": null, "metadata": {"description": "SciFact (MC) using SciLLM", "regimes": ["SciLLM-v0.1"], "alias": "scifact:mc::scillm"}}, "current_date": "2024-10-13 03:46:14 UTC", "num_instances": 340}

src/eval/oe_eval_tasks/scifact/rc_3shot/config.json

@@ -0,0 +1 @@
+{"task_name": "scifact", "task_hash": "d5ef8896b403d58b663afe1be9ea4a62", "task_config": {"task_name": "scifact", "task_core": "scifact", "limit": 1000, "split": "validation", "num_shots": 3, "fewshot_seed": 1234, "primary_metric": "acc_raw", "random_subsample_seed": 1234, "context_kwargs": {"description": null}, "generation_kwargs": {}, "metric_kwargs": {"uncond_docid_offset": 1000000}, "native_id_field": "id", "fewshot_source": "SciLLM:SciFact", "dataset_path": "allenai/scifact_entailment", "dataset_name": null, "use_chat_format": null, "version": 0, "revision": null, "metadata": {"description": "SciFact (RC) using SciLLM", "regimes": ["SciLLM-v0.1"], "alias": "scifact:rc::scillm"}}, "current_date": "2024-10-13 03:46:06 UTC", "num_instances": 340}

src/examples/train.py

@@ -34,6 +34,7 @@
     GPUMemoryMonitorCallback,
     GradClipperCallback,
     LMEvaluatorCallbackConfig,
+    DownstreamEvaluatorCallbackConfig,
     ProfilerCallback,
     SchedulerCallback,
     SequenceLengthSchedulerCallback,
@@ -89,6 +90,11 @@ def build_config(run_name: str, overrides: List[str]) -> ExperimentConfig:
         num_workers=4,
     )
 
+    downstream_evaluators = [
+        "pubmedqa_mc",
+        "scifact_rc",
+    ]
+
     trainer_config = (
         TrainerConfig(
             save_folder=f"/tmp/{run_name}",
@@ -147,6 +153,16 @@ def build_config(run_name: str, overrides: List[str]) -> ExperimentConfig:
                 eval_duration=Duration.steps(10),
             ),
         )
+        .with_callback(
+            "downstream",
+            DownstreamEvaluatorCallbackConfig(
+                labels=downstream_evaluators,
+                eval_batch_size=4,
+                tokenizer="Qwen/Qwen2.5-1.5B",
+                eval_interval=250,
+                # eval_duration=Duration.steps(10),
+            ),
+        )
     )
 
     return ExperimentConfig(

src/olmo_core/train/callbacks/__init__.py

@@ -3,6 +3,7 @@
 from .comet import CometCallback, CometNotificationSetting
 from .config_saver import ConfigSaverCallback
 from .console_logger import ConsoleLoggerCallback
+from .downstream import DownstreamEvaluatorCallbackConfig
 from .evaluator_callback import EvaluatorCallback, LMEvaluatorCallbackConfig
 from .float8_handler import Float8HandlerCallback
 from .garbage_collector import GarbageCollectorCallback
@@ -24,6 +25,7 @@
     "CometNotificationSetting",
     "ConfigSaverCallback",
     "ConsoleLoggerCallback",
+    "DownstreamEvaluatorCallbackConfig",
     "EvaluatorCallback",
     "Float8HandlerCallback",
     "LMEvaluatorCallbackConfig",