From bab852613e00734049a59dea4448fdd0a76107c5 Mon Sep 17 00:00:00 2001 From: "Peter B. Johnson" Date: Wed, 4 Nov 2026 22:00:24 +0000 Subject: [PATCH 1/2] Extract shared Hugging Face model code from gpt2_infer Pure refactor in preparation for a second Hugging Face model; GPT-2 output is unchanged (verified identical for fixed seeds). - models/hf_lm.py: cached local-dir-or-download loading, sampling, param clamping and error result, previously private to gpt2_infer - Load weights explicitly in float32 rather than relying on the transformers default - Rename scripts/download_gpt2.py to download_hf_model.py, taking the target dir and model id explicitly Co-Authored-By: Claude Opus 5.5 --- Dockerfile | 6 +- docs/dev.md | 2 +- evaluation_function/evaluation_test.py | 2 +- evaluation_function/models/gpt2_infer.py | 118 ++++------------------ evaluation_function/models/hf_lm.py | 120 +++++++++++++++++++++++ scripts/download_gpt2.py | 46 --------- scripts/download_hf_model.py | 36 +++++++ 7 files changed, 182 insertions(+), 148 deletions(-) create mode 100644 evaluation_function/models/hf_lm.py delete mode 100644 scripts/download_gpt2.py create mode 100644 scripts/download_hf_model.py diff --git a/Dockerfile b/Dockerfile index 385ed4a..6c3ac7e 100755 --- a/Dockerfile +++ b/Dockerfile @@ -35,13 +35,13 @@ COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV} # /tmp, the only writable path on Lambda. ENV GPT2_MODEL_DIR=/app/hf_models/gpt2 \ HF_HOME=/tmp/hf -COPY scripts/download_gpt2.py /tmp/download_gpt2.py +COPY scripts/download_hf_model.py /tmp/download_hf_model.py # safetensors >=0.8 writes model.safetensors via a tempfile+rename, leaving it # mode 0600 (root-only). Lambda runs as a non-root user, so open() fails (and # safetensors reports that as "No such file or directory"). Make it readable. -RUN HF_HOME=/tmp/hf_build python /tmp/download_gpt2.py "${GPT2_MODEL_DIR}" && \ +RUN HF_HOME=/tmp/hf_build python /tmp/download_hf_model.py "${GPT2_MODEL_DIR}" --model-id gpt2 && \ chmod -R a+rX "${GPT2_MODEL_DIR}" && \ - rm -rf /tmp/hf_build /tmp/download_gpt2.py + rm -rf /tmp/hf_build /tmp/download_hf_model.py # Copy evaluation function first (smaller, changes more often) COPY evaluation_function ./evaluation_function diff --git a/docs/dev.md b/docs/dev.md index 5ed1440..5fcae62 100644 --- a/docs/dev.md +++ b/docs/dev.md @@ -11,7 +11,7 @@ The parameters for each model are documented for teachers in `user.md`. ## Models: build and runtime notes ### `gpt2_infer` -- `scripts/download_gpt2.py` downloads the weights into the image during `docker build`. They go to `GPT2_MODEL_DIR` (`/app/hf_models/gpt2`), so nothing is fetched at runtime. +- `scripts/download_hf_model.py` downloads the weights into the image during `docker build`. They go to `GPT2_MODEL_DIR` (`/app/hf_models/gpt2`), so nothing is fetched at runtime. - Locally, the weights download on first use into `HF_HOME` (default `/tmp/hf`). - The test that does a real download and generation is skipped unless `RUN_GPT2_TESTS=1` is set, so CI doesn't download about 500MB. - Memory: see the note above (2048MB minimum). diff --git a/evaluation_function/evaluation_test.py b/evaluation_function/evaluation_test.py index 348f278..f167dc0 100755 --- a/evaluation_function/evaluation_test.py +++ b/evaluation_function/evaluation_test.py @@ -57,7 +57,7 @@ def test_generates_continuation(self): def test_load_failure_returns_error_result(self): from .models import gpt2_infer - with mock.patch.object(gpt2_infer, "_load", side_effect=OSError("no weights")): + with mock.patch.object(gpt2_infer.MODEL, "load", side_effect=OSError("no weights")): result = evaluation_function( "hello", None, {"model": "gpt2_infer"} ).to_dict() diff --git a/evaluation_function/models/gpt2_infer.py b/evaluation_function/models/gpt2_infer.py index 483d40f..680caf6 100644 --- a/evaluation_function/models/gpt2_infer.py +++ b/evaluation_function/models/gpt2_infer.py @@ -3,12 +3,10 @@ The next rung on the ladder after the Bengio neural n-gram model: a full transformer language model. The student's response is used as the prompt and -the model's sampled continuation is returned as feedback. +the model's sampled continuation (prompt included) is returned as feedback. -Weights are baked into the Docker image at build time (see ``Dockerfile``) and -located via the ``GPT2_MODEL_DIR`` env var. When that directory is missing -(e.g. local development) the model is downloaded from Hugging Face on first use -into ``HF_HOME`` (default ``/tmp/hf``, the only writable path on Lambda). +Weights come from ``GPT2_MODEL_DIR`` when present, otherwise they are downloaded +on first use (see ``hf_lm``). Params: - ``word_count``: number of new tokens to generate (default 20, max 100) @@ -16,114 +14,40 @@ - ``top_k``: sample only from the k most likely tokens (default 50) - ``seed``: optional integer for reproducible output """ -import os -import traceback -from pathlib import Path - from lf_toolkit.evaluation import Result, Params -from evaluation_function.lazy_load import LazyModule - -torch = LazyModule("torch") -transformers = LazyModule("transformers") +from evaluation_function.models.hf_lm import ( + HFCausalLM, + failure_result, + prompt_text, + sample, + sampling_params, +) -DEFAULT_MODEL_ID = "gpt2" +MODEL = HFCausalLM("GPT-2", "gpt2", dir_env="GPT2_MODEL_DIR", id_env="GPT2_MODEL_ID") DEFAULT_PROMPT = "the general" DEFAULT_TOKENS = 20 MAX_TOKENS = 100 - -# Cached across requests so a warm worker loads the weights only once. -_MODEL = None -_TOKENIZER = None -_DEVICE = None - - -def _load(): - global _MODEL, _TOKENIZER, _DEVICE - if _MODEL is not None: - return _MODEL, _TOKENIZER, _DEVICE - - model_dir = os.environ.get("GPT2_MODEL_DIR") - if model_dir and Path(model_dir).is_dir(): - print(f"### Loading GPT-2 from {model_dir} ###") - source, kwargs = model_dir, {"local_files_only": True} - else: - source = os.environ.get("GPT2_MODEL_ID", DEFAULT_MODEL_ID) - cache_dir = os.environ.get("HF_HOME", "/tmp/hf") - print(f"### Downloading GPT-2 '{source}' into {cache_dir} ###") - kwargs = {"cache_dir": cache_dir} - - tokenizer = transformers.AutoTokenizer.from_pretrained(source, **kwargs) - model = transformers.AutoModelForCausalLM.from_pretrained(source, **kwargs) - - device = torch.device("mps" if torch.backends.mps.is_available() else "cpu") - model.to(device).eval() - - _MODEL, _TOKENIZER, _DEVICE = model, tokenizer, device - return _MODEL, _TOKENIZER, _DEVICE - - -def _int_param(params, name, default, lo, hi): - try: - value = int(params.get(name, default)) - except (TypeError, ValueError): - value = default - return max(lo, min(hi, value)) - - -def _float_param(params, name, default, lo, hi): - try: - value = float(params.get(name, default)) - except (TypeError, ValueError): - value = default - return max(lo, min(hi, value)) +VOCAB_SIZE = 50257 def complete(prompt, max_new_tokens=DEFAULT_TOKENS, temperature=1.0, top_k=50, seed=None): - model, tokenizer, device = _load() - if seed is not None: - torch.manual_seed(int(seed)) - - inputs = tokenizer(prompt, return_tensors="pt").to(device) - with torch.no_grad(): - output = model.generate( - **inputs, - max_new_tokens=max_new_tokens, - do_sample=True, - temperature=temperature, - top_k=top_k, - pad_token_id=tokenizer.eos_token_id, - ) - return tokenizer.decode(output[0], skip_special_tokens=True) + model, tokenizer, device = MODEL.load() + input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device) + output = sample(model, tokenizer, input_ids, max_new_tokens, temperature, top_k, seed) + return tokenizer.decode(output, skip_special_tokens=True) def run(response, answer, params: Params) -> Result: # Guard the whole load+infer path so a failed download or corrupt weights # degrade this one request rather than taking down the RPC worker. try: - return _run_inference(response, params) - except Exception as e: - print("### gpt2_infer failed ###") - traceback.print_exc() - return Result( - is_correct=False, - feedback_items=[ - ("general", "Could not run the GPT-2 language model."), - ("error", str(e)), - ], + text = complete( + prompt_text(response, DEFAULT_PROMPT), + **sampling_params(params, DEFAULT_TOKENS, MAX_TOKENS, VOCAB_SIZE), ) - - -def _run_inference(response, params: Params) -> Result: - prompt = response.strip() if isinstance(response, str) and response.strip() else DEFAULT_PROMPT - - text = complete( - prompt, - max_new_tokens=_int_param(params, "word_count", DEFAULT_TOKENS, 1, MAX_TOKENS), - temperature=_float_param(params, "temperature", 1.0, 0.01, 5.0), - top_k=_int_param(params, "top_k", 50, 1, 50257), - seed=params.get("seed"), - ) + except Exception as e: + return failure_result("gpt2_infer", "GPT-2", e) result = Result(True) result.add_feedback("general", text) diff --git a/evaluation_function/models/hf_lm.py b/evaluation_function/models/hf_lm.py new file mode 100644 index 0000000..68dbc40 --- /dev/null +++ b/evaluation_function/models/hf_lm.py @@ -0,0 +1,120 @@ +"""Shared loading and sampling for the pre-trained Hugging Face language models. + +Weights are baked into the Docker image at build time (see ``Dockerfile``) and +located via a per-model ``*_MODEL_DIR`` env var. When that directory is missing +(e.g. local development) the model is downloaded from Hugging Face on first use +into ``HF_HOME`` (default ``/tmp/hf``, the only writable path on Lambda). + +Models are loaded in float32: Lambda is CPU-only, where half-precision is slow +and quantisation would distort the comparison between models. +""" +import os +import traceback +from pathlib import Path + +from lf_toolkit.evaluation import Result + +from evaluation_function.lazy_load import LazyModule + +torch = LazyModule("torch") +transformers = LazyModule("transformers") + + +class HFCausalLM: + """A pre-trained causal LM, loaded on first use and cached for the worker's lifetime.""" + + def __init__(self, label, default_model_id, dir_env, id_env): + self.label = label + self.default_model_id = default_model_id + self.dir_env = dir_env + self.id_env = id_env + self._loaded = None + + def load(self): + """Return ``(model, tokenizer, device)``.""" + if self._loaded is None: + self._loaded = self._load() + return self._loaded + + def _load(self): + model_dir = os.environ.get(self.dir_env) + if model_dir and Path(model_dir).is_dir(): + print(f"### Loading {self.label} from {model_dir} ###") + source, kwargs = model_dir, {"local_files_only": True} + else: + source = os.environ.get(self.id_env, self.default_model_id) + cache_dir = os.environ.get("HF_HOME", "/tmp/hf") + print(f"### Downloading {self.label} '{source}' into {cache_dir} ###") + kwargs = {"cache_dir": cache_dir} + + tokenizer = transformers.AutoTokenizer.from_pretrained(source, **kwargs) + model = transformers.AutoModelForCausalLM.from_pretrained(source, dtype=torch.float32, **kwargs) + + device = torch.device("mps" if torch.backends.mps.is_available() else "cpu") + model.to(device).eval() + return model, tokenizer, device + + +def sample(model, tokenizer, input_ids, max_new_tokens, temperature, top_k, seed=None): + """Sample a continuation of ``input_ids``; returns the full token sequence (prompt included).""" + if seed is not None: + torch.manual_seed(int(seed)) + + pad_token_id = tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = tokenizer.eos_token_id + + with torch.no_grad(): + output = model.generate( + input_ids, + attention_mask=torch.ones_like(input_ids), + max_new_tokens=max_new_tokens, + do_sample=True, + temperature=temperature, + top_k=top_k, + pad_token_id=pad_token_id, + ) + return output[0] + + +def sampling_params(params, default_tokens, max_tokens, vocab_size): + """Read the sampling params shared by all the Hugging Face models, clamped to valid ranges.""" + return { + "max_new_tokens": _int_param(params, "word_count", default_tokens, 1, max_tokens), + "temperature": _float_param(params, "temperature", 1.0, 0.01, 5.0), + "top_k": _int_param(params, "top_k", 50, 1, vocab_size), + "seed": params.get("seed"), + } + + +def prompt_text(response, default): + return response.strip() if isinstance(response, str) and response.strip() else default + + +def failure_result(name, label, exc) -> Result: + """Turn a load or inference failure into a feedback result rather than crashing the worker.""" + print(f"### {name} failed ###") + traceback.print_exc() + return Result( + is_correct=False, + feedback_items=[ + ("general", f"Could not run the {label} language model."), + ("error", str(exc)), + ], + ) + + +def _int_param(params, name, default, lo, hi): + try: + value = int(params.get(name, default)) + except (TypeError, ValueError): + value = default + return max(lo, min(hi, value)) + + +def _float_param(params, name, default, lo, hi): + try: + value = float(params.get(name, default)) + except (TypeError, ValueError): + value = default + return max(lo, min(hi, value)) diff --git a/scripts/download_gpt2.py b/scripts/download_gpt2.py deleted file mode 100644 index 70486ca..0000000 --- a/scripts/download_gpt2.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Download the pre-trained GPT-2 weights and tokenizer into a local directory. - -Used at Docker build time to bake the model into the image so that -``evaluation_function.models.gpt2_infer`` needs no network access at runtime. - -Usage: - python scripts/download_gpt2.py [TARGET_DIR] [--model-id gpt2] - -TARGET_DIR defaults to $GPT2_MODEL_DIR. -""" -import argparse -import os -import sys - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) - parser.add_argument( - "target_dir", - nargs="?", - default=os.environ.get("GPT2_MODEL_DIR"), - help="Directory to save the model into (default: $GPT2_MODEL_DIR)", - ) - parser.add_argument( - "--model-id", - default=os.environ.get("GPT2_MODEL_ID", "gpt2"), - help="Hugging Face model id (default: $GPT2_MODEL_ID or 'gpt2')", - ) - args = parser.parse_args() - - if not args.target_dir: - parser.error("no target directory given and GPT2_MODEL_DIR is not set") - - from transformers import AutoModelForCausalLM, AutoTokenizer - - print(f"Downloading '{args.model_id}' into {args.target_dir}") - AutoTokenizer.from_pretrained(args.model_id).save_pretrained(args.target_dir) - AutoModelForCausalLM.from_pretrained( - args.model_id, use_safetensors=True - ).save_pretrained(args.target_dir) - print("Done") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/scripts/download_hf_model.py b/scripts/download_hf_model.py new file mode 100644 index 0000000..540720c --- /dev/null +++ b/scripts/download_hf_model.py @@ -0,0 +1,36 @@ +"""Download a pre-trained Hugging Face causal LM and its tokenizer into a local directory. + +Used at Docker build time to bake the models into the image so that the +Hugging Face models (see ``evaluation_function/models/hf_lm.py``) need no +network access at runtime. Weights are saved in float32, the precision they are served in. + +Usage: + python scripts/download_hf_model.py TARGET_DIR --model-id MODEL_ID + +e.g. + python scripts/download_hf_model.py /app/hf_models/gpt2 --model-id gpt2 +""" +import argparse +import sys + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("target_dir", help="Directory to save the model into") + parser.add_argument("--model-id", required=True, help="Hugging Face model id, e.g. 'gpt2'") + args = parser.parse_args() + + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + + print(f"Downloading '{args.model_id}' into {args.target_dir}") + AutoTokenizer.from_pretrained(args.model_id).save_pretrained(args.target_dir) + AutoModelForCausalLM.from_pretrained( + args.model_id, dtype=torch.float32, use_safetensors=True + ).save_pretrained(args.target_dir) + print("Done") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 6cb0674f798deccd680c4317e45630375fe116ff Mon Sep 17 00:00:00 2001 From: "Peter B. Johnson" Date: Wed, 4 Nov 2026 22:08:42 +0000 Subject: [PATCH 2/2] Add SmolLM2-135M-Instruct model (smollm2_infer) Adds SmolLM2-135M-Instruct alongside GPT-2 for side-by-side comparison in the course: similar parameter count, but newer data, architecture and instruction tuning. GPT-2 is unchanged. - smollm2_infer wraps the student's response in the chat template and returns only the model's reply (GPT-2 returns prompt + continuation); the difference is documented for teachers in docs/user.md - Same params as gpt2_infer; word_count defaults to 64 (max 256) for chat replies and top_k is clamped to the 49152-token vocabulary - Bake the weights into the image under SMOLLM2_MODEL_DIR, in their own layer, with the same safetensors permission fix as GPT-2 - Tests cover both models; generation tests now run under RUN_HF_MODEL_TESTS (was RUN_GPT2_TESTS) Verified in the built image (offline, non-root): both models load from the baked directories, with peak memory about 1.5GB once both are loaded. Co-Authored-By: Claude Opus 5.5 --- Dockerfile | 13 +++- docs/dev.md | 13 ++-- docs/user.md | 10 +++ evaluation_function/dev.json | 8 +++ evaluation_function/evaluation_test.py | 45 ++++++++++---- evaluation_function/models/__init__.py | 1 + evaluation_function/models/smollm2_infer.py | 68 +++++++++++++++++++++ 7 files changed, 136 insertions(+), 22 deletions(-) create mode 100644 evaluation_function/models/smollm2_infer.py diff --git a/Dockerfile b/Dockerfile index 6c3ac7e..ecc95b8 100755 --- a/Dockerfile +++ b/Dockerfile @@ -30,17 +30,24 @@ ENV VIRTUAL_ENV=/app/.venv \ # Copy the cleaned virtual environment COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV} -# Bake the pre-trained GPT-2 weights into the image so gpt2_infer needs no -# network access at runtime. HF_HOME points any runtime fallback download at -# /tmp, the only writable path on Lambda. +# Bake the pre-trained Hugging Face weights (GPT-2 and SmolLM2, ~540MB each in +# float32) into the image so gpt2_infer and smollm2_infer need no network +# access at runtime. HF_HOME points any runtime fallback download at /tmp, the +# only writable path on Lambda. ENV GPT2_MODEL_DIR=/app/hf_models/gpt2 \ + SMOLLM2_MODEL_DIR=/app/hf_models/smollm2-135m-instruct \ HF_HOME=/tmp/hf COPY scripts/download_hf_model.py /tmp/download_hf_model.py # safetensors >=0.8 writes model.safetensors via a tempfile+rename, leaving it # mode 0600 (root-only). Lambda runs as a non-root user, so open() fails (and # safetensors reports that as "No such file or directory"). Make it readable. +# One layer per model, so changing one does not re-download the other. RUN HF_HOME=/tmp/hf_build python /tmp/download_hf_model.py "${GPT2_MODEL_DIR}" --model-id gpt2 && \ chmod -R a+rX "${GPT2_MODEL_DIR}" && \ + rm -rf /tmp/hf_build +RUN HF_HOME=/tmp/hf_build python /tmp/download_hf_model.py "${SMOLLM2_MODEL_DIR}" \ + --model-id HuggingFaceTB/SmolLM2-135M-Instruct && \ + chmod -R a+rX "${SMOLLM2_MODEL_DIR}" && \ rm -rf /tmp/hf_build /tmp/download_hf_model.py # Copy evaluation function first (smaller, changes more often) diff --git a/docs/dev.md b/docs/dev.md index 5fcae62..9a3975d 100644 --- a/docs/dev.md +++ b/docs/dev.md @@ -1,7 +1,7 @@ # langModels An evaluation function to aid the learning of language models. -**NOTE:** AWS deployments need at least 2048MB RAM, and storage to match the ~3.4GB image. GPT-2 (`gpt2_infer`) peaks near 1GB in the Python process alone, so 1024MB is no longer enough. +**NOTE:** AWS deployments need at least 2048MB RAM, and storage to match the ~4GB image. GPT-2 (`gpt2_infer`) peaks near 1GB in the Python process alone, so 1024MB is no longer enough. A warm worker that has served both GPT-2 and SmolLM2 (`smollm2_infer`) keeps both loaded: peak memory measured in the container was about 1.5GB, which leaves around 500MB of headroom at 2048MB. ## Inputs *Specific input parameters which can be supplied when the `eval` command is supplied to this function.* @@ -10,11 +10,12 @@ The parameters for each model are documented for teachers in `user.md`. ## Models: build and runtime notes -### `gpt2_infer` -- `scripts/download_hf_model.py` downloads the weights into the image during `docker build`. They go to `GPT2_MODEL_DIR` (`/app/hf_models/gpt2`), so nothing is fetched at runtime. -- Locally, the weights download on first use into `HF_HOME` (default `/tmp/hf`). -- The test that does a real download and generation is skipped unless `RUN_GPT2_TESTS=1` is set, so CI doesn't download about 500MB. -- Memory: see the note above (2048MB minimum). +### Hugging Face models: `gpt2_infer` and `smollm2_infer` +Both share their loading and sampling code (`evaluation_function/models/hf_lm.py`), and are loaded in float32 with no quantisation, since Lambda runs on CPU. +- `scripts/download_hf_model.py` downloads the weights into the image during `docker build`, one layer per model. They go to `GPT2_MODEL_DIR` (`/app/hf_models/gpt2`) and `SMOLLM2_MODEL_DIR` (`/app/hf_models/smollm2-135m-instruct`), so nothing is fetched at runtime. +- Locally, the weights download on first use into `HF_HOME` (default `/tmp/hf`). Set `GPT2_MODEL_DIR` or `SMOLLM2_MODEL_DIR` to use a directory made by the download script instead. +- The tests that do a real download and generation are skipped unless `RUN_HF_MODEL_TESTS=1` is set, so CI doesn't download about 1GB. +- Memory: see the note above. ## Outputs *Output schema/values for this function* diff --git a/docs/user.md b/docs/user.md index 305bd6f..1fcb315 100644 --- a/docs/user.md +++ b/docs/user.md @@ -15,3 +15,13 @@ Pre-trained GPT-2 (124M) from Hugging Face. The student's response is used as th | `temperature` | 1.0 | Sampling temperature (0.01–5.0): higher values give more varied text | | `top_k` | 50 | Sample only from the k most likely next tokens (1–50257) | | `seed` | none | Integer for reproducible output | + +### `smollm2_infer` +Pre-trained SmolLM2-135M-Instruct from Hugging Face, for comparison with GPT-2: a similar size (135M parameters), but trained on newer data, with a newer architecture, and instruction-tuned to act as a chat assistant. The student's response is sent as a chat message, and only the model's reply is returned. Unlike GPT-2, the prompt is not repeated at the start of the feedback. The model's chat template adds a default system prompt ("You are a helpful AI assistant named SmolLM, trained by Hugging Face"), so replies may refer to the model as SmolLM. + +| Param | Default | Notes | +|---|---|---| +| `word_count` | 64 | Maximum new tokens to generate (1–256). The reply can be shorter if the model finishes its answer | +| `temperature` | 1.0 | Sampling temperature (0.01–5.0): higher values give more varied text | +| `top_k` | 50 | Sample only from the k most likely next tokens (1–49152) | +| `seed` | none | Integer for reproducible output | diff --git a/evaluation_function/dev.json b/evaluation_function/dev.json index f331042..29db868 100644 --- a/evaluation_function/dev.json +++ b/evaluation_function/dev.json @@ -58,5 +58,13 @@ "word_count": 20, "temperature": 1.0, "top_k": 50 + }, + "smollm2_infer": { + "response": "What is a language model?", + "answer": 0, + "model": "smollm2_infer", + "word_count": 64, + "temperature": 1.0, + "top_k": 50 } } diff --git a/evaluation_function/evaluation_test.py b/evaluation_function/evaluation_test.py index f167dc0..51a2e23 100755 --- a/evaluation_function/evaluation_test.py +++ b/evaluation_function/evaluation_test.py @@ -40,12 +40,12 @@ def test_evaluation(self): self.assertIsInstance(feedback, str) -class TestGPT2(unittest.TestCase): - """GPT-2 tests. The happy path needs the weights (baked into the image, or a - ~500MB Hugging Face download), so it only runs when RUN_GPT2_TESTS is set.""" +class TestHFModels(unittest.TestCase): + """GPT-2 and SmolLM2 tests. Generation needs the weights (baked into the image, + or ~500MB Hugging Face downloads each), so it only runs when RUN_HF_MODEL_TESTS is set.""" - @unittest.skipUnless(os.environ.get("RUN_GPT2_TESTS"), "set RUN_GPT2_TESTS=1 to run") - def test_generates_continuation(self): + @unittest.skipUnless(os.environ.get("RUN_HF_MODEL_TESTS"), "set RUN_HF_MODEL_TESTS=1 to run") + def test_gpt2_continues_prompt(self): prompt = "Once upon a time" result = evaluation_function( prompt, None, {"model": "gpt2_infer", "word_count": 5, "seed": 0} @@ -54,12 +54,31 @@ def test_generates_continuation(self): self.assertTrue(result["feedback"].startswith(prompt)) self.assertGreater(len(result["feedback"]), len(prompt)) + @unittest.skipUnless(os.environ.get("RUN_HF_MODEL_TESTS"), "set RUN_HF_MODEL_TESTS=1 to run") + def test_smollm2_replies_without_prompt(self): + prompt = "What is the capital of France?" + result = evaluation_function( + prompt, None, {"model": "smollm2_infer", "word_count": 20, "seed": 0} + ).to_dict() + self.assertTrue(result.get("is_correct")) + self.assertTrue(result["feedback"]) + self.assertNotIn(prompt, result["feedback"]) + self.assertNotIn("<|im_", result["feedback"]) + + def test_top_k_clamped_to_vocab(self): + from .models import gpt2_infer, smollm2_infer + from .models.hf_lm import sampling_params + + for module, vocab_size in [(gpt2_infer, 50257), (smollm2_infer, 49152)]: + sampling = sampling_params({"top_k": 10**9}, module.DEFAULT_TOKENS, module.MAX_TOKENS, module.VOCAB_SIZE) + self.assertEqual(sampling["top_k"], vocab_size) + def test_load_failure_returns_error_result(self): - from .models import gpt2_infer - - with mock.patch.object(gpt2_infer.MODEL, "load", side_effect=OSError("no weights")): - result = evaluation_function( - "hello", None, {"model": "gpt2_infer"} - ).to_dict() - self.assertFalse(result.get("is_correct")) - self.assertIn("no weights", str(result.get("feedback"))) \ No newline at end of file + from .models import gpt2_infer, smollm2_infer + + for name, module in [("gpt2_infer", gpt2_infer), ("smollm2_infer", smollm2_infer)]: + with self.subTest(model=name), \ + mock.patch.object(module.MODEL, "load", side_effect=OSError("no weights")): + result = evaluation_function("hello", None, {"model": name}).to_dict() + self.assertFalse(result.get("is_correct")) + self.assertIn("no weights", str(result.get("feedback"))) diff --git a/evaluation_function/models/__init__.py b/evaluation_function/models/__init__.py index 454cc2b..78d2add 100644 --- a/evaluation_function/models/__init__.py +++ b/evaluation_function/models/__init__.py @@ -12,6 +12,7 @@ "shannon_words_ngram", "bengio_infer", "gpt2_infer", + "smollm2_infer", ] __all__ = ["AVAILABLE_MODELS"] \ No newline at end of file diff --git a/evaluation_function/models/smollm2_infer.py b/evaluation_function/models/smollm2_infer.py new file mode 100644 index 0000000..ee61fdf --- /dev/null +++ b/evaluation_function/models/smollm2_infer.py @@ -0,0 +1,68 @@ +""" +SmolLM2-135M-Instruct inference using the pre-trained Hugging Face checkpoint. + +Shown alongside ``gpt2_infer`` for comparison: a similar parameter count, but +newer training data, a newer (Llama-style) architecture and instruction tuning. +The student's response is wrapped in the model's chat template as a user +message, and only the model's reply is returned as feedback. Unlike GPT-2, the +prompt is not echoed: the student asked a question rather than starting a +sentence for the model to finish. + +Weights come from ``SMOLLM2_MODEL_DIR`` when present, otherwise they are +downloaded on first use (see ``hf_lm``). + +Params: +- ``word_count``: maximum number of new tokens to generate (default 64, max 256); + the reply may end sooner when the model finishes its turn +- ``temperature``: sampling temperature (default 1.0) +- ``top_k``: sample only from the k most likely tokens (default 50) +- ``seed``: optional integer for reproducible output +""" +from lf_toolkit.evaluation import Result, Params + +from evaluation_function.models.hf_lm import ( + HFCausalLM, + failure_result, + prompt_text, + sample, + sampling_params, +) + +MODEL = HFCausalLM( + "SmolLM2", + "HuggingFaceTB/SmolLM2-135M-Instruct", + dir_env="SMOLLM2_MODEL_DIR", + id_env="SMOLLM2_MODEL_ID", +) +DEFAULT_PROMPT = "What is a language model?" +DEFAULT_TOKENS = 64 +MAX_TOKENS = 256 +VOCAB_SIZE = 49152 + + +def reply(prompt, max_new_tokens=DEFAULT_TOKENS, temperature=1.0, top_k=50, seed=None): + model, tokenizer, device = MODEL.load() + input_ids = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + return_dict=True, + return_tensors="pt", + )["input_ids"].to(device) + output = sample(model, tokenizer, input_ids, max_new_tokens, temperature, top_k, seed) + return tokenizer.decode(output[input_ids.shape[-1]:], skip_special_tokens=True).strip() + + +def run(response, answer, params: Params) -> Result: + # Guard the whole load+infer path so a failed download or corrupt weights + # degrade this one request rather than taking down the RPC worker. + try: + text = reply( + prompt_text(response, DEFAULT_PROMPT), + **sampling_params(params, DEFAULT_TOKENS, MAX_TOKENS, VOCAB_SIZE), + ) + except Exception as e: + return failure_result("smollm2_infer", "SmolLM2", e) + + result = Result(True) + result.add_feedback("general", text) + return result