diff --git a/README.md b/README.md index a0885cdd..4d12f592 100644 --- a/README.md +++ b/README.md @@ -331,6 +331,17 @@ env = MyTeacherEnv( ) ``` +You can either: + +- build a teacher-enabled env by mixing `TeacherDistillationEnv` into an existing + `BaseEnv`-derived env such as `GSM8kEnv`, or +- subclass `TeacherDistillationEnv` directly and implement the usual environment + methods yourself. + +In both cases, `TeacherDistillationEnv` still assumes the normal `BaseEnv` +runtime contract: tokenized rollouts, `ScoredDataGroup` payloads, and the +standard `handle_send_to_api(...)` transport path. + CLI shape: ```bash @@ -345,6 +356,13 @@ If `--teacher.model_name` is a deployment alias rather than a tokenizer identifier, also set `--teacher.tokenizer_name ...` so the env can validate tokenizer compatibility. +Scope note: + +- The teacher-aware CLI wiring currently exists for `serve`. +- If `teacher_enabled=True`, the generic `process` and `evaluate` commands will + fail loudly at env construction time unless you instantiate the env yourself + and pass `teacher_server_configs=...`. + Tokenizer requirement: - Teacher distillation currently requires the teacher and student to use the same tokenizer vocabulary. diff --git a/atroposlib/envs/teacher_distillation_env.py b/atroposlib/envs/teacher_distillation_env.py index 892d424c..64f62b14 100644 --- a/atroposlib/envs/teacher_distillation_env.py +++ b/atroposlib/envs/teacher_distillation_env.py @@ -296,7 +296,11 @@ class TeacherDistillationEnv(BaseEnv, ABC): if config.teacher_enabled: if teacher_server_configs is None: raise ValueError( - "teacher_enabled=True requires teacher_server_configs at init." + "teacher_enabled=True but no teacher server configuration was " + "provided. Pass teacher_server_configs=... when instantiating " + "the environment directly, or use the teacher-aware 'serve' CLI " + "path with --teacher.* flags. The generic BaseEnv 'process' and " + "'evaluate' commands do not currently wire teacher_server_configs." ) teacher_config_source = teacher_server_configs self.teacher_server = ServerManager( diff --git a/atroposlib/tests/test_teacher_distillation_env.py b/atroposlib/tests/test_teacher_distillation_env.py index b348e65a..b35c48ee 100644 --- a/atroposlib/tests/test_teacher_distillation_env.py +++ b/atroposlib/tests/test_teacher_distillation_env.py @@ -205,7 +205,7 @@ def test_init_requires_teacher_server_source(monkeypatch): teacher_enabled=True, teacher_top_k=0, ) - with pytest.raises(ValueError, match="teacher_enabled=True requires"): + with pytest.raises(ValueError, match="no teacher server configuration was provided"): _ConcreteTeacherEnv( config=config, server_configs=[], diff --git a/example_trainer/README.md b/example_trainer/README.md index de31b3eb..4563e12c 100644 --- a/example_trainer/README.md +++ b/example_trainer/README.md @@ -324,6 +324,11 @@ If `$TEACHER_MODEL` is a deployment alias instead of a tokenizer identifier, also set `--teacher.tokenizer_name ...` so the env can validate tokenizer compatibility. +The teacher-aware CLI path is currently wired for `serve`. If +`teacher_enabled=True`, the generic `process` and `evaluate` commands are not +teacher-aware and will fail loudly unless the environment is instantiated +manually with `teacher_server_configs=...`. + Why cross-tokenizer conversion is not acceptable here: - Teacher token ID `1234` and student token ID `1234` can correspond to different text.