Add IS evaluator

fast_llm/engine/evaluation/config.py

@@ -8,6 +8,7 @@
 
 if typing.TYPE_CHECKING:
     from fast_llm.engine.evaluation.evaluator import Evaluator, EvaluatorLmEval, LossEvaluator
+    from fast_llm.engine.evaluation.forward_kl.evaluator import ForwardKLEvaluator
 
 
 @config_class()
@@ -119,3 +120,52 @@ def get_evaluator(
         from fast_llm.engine.evaluation.lm_eval.evaluator import LmEvalEvaluator
 
         return LmEvalEvaluator(name, self, batch_config, data_load_num_proc, train_iters)
+
+
+@config_class(dynamic_type={EvaluatorConfig: "forward_kl"})
+class ForwardKLEvaluatorConfig(EvaluatorConfig):
+    _abstract: typing.ClassVar[bool] = False
+
+    dataset_path: str | None = Field(
+        default=None,
+        desc="HuggingFace dataset path containing teacher traces.",
+        hint=FieldHint.core,
+    )
+    split: str = Field(
+        default="validation",
+        desc="Dataset split to evaluate on. Use 'train+validation' syntax to combine multiple splits.",
+        hint=FieldHint.optional,
+    )
+    seed: int = Field(
+        default=42,
+        desc="Random seed for shuffling traces. Ensures reproducible evaluation across runs.",
+        hint=FieldHint.optional,
+    )
+    num_samples: int | None = Field(
+        default=None,
+        desc="Maximum number of traces to evaluate (after shuffling). None for all.",
+        hint=FieldHint.optional,
+        valid=skip_valid_if_none(check_field(Assert.gt, 0)),
+    )
+    batch_size: int = Field(
+        default=8,
+        desc="Batch size for forward passes.",
+        hint=FieldHint.performance,
+        valid=check_field(Assert.gt, 0),
+    )
+    trust_remote_code: bool = Field(
+        default=False,
+        desc="Trust remote code when loading dataset.",
+        hint=FieldHint.optional,
+    )
+
+    def get_evaluator(
+        self,
+        name: str,
+        batch_config: BatchConfig,
+        data_load_num_proc: int,
+        train_iters: int | None = None,
+    ) -> "ForwardKLEvaluator":
+        from fast_llm.engine.evaluation.forward_kl.evaluator import ForwardKLEvaluator
+
+        return ForwardKLEvaluator(name, self, batch_config, data_load_num_proc, train_iters)