Add SmolLM2-135M-Instruct model (smollm2_infer)

Dockerfile

@@ -30,18 +30,25 @@ 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_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}" && \
+# 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 /tmp/download_gpt2.py
+    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)
 COPY evaluation_function ./evaluation_function

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_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.
-- 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*

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 |

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
   }
 }

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, "_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")))
+        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")))

evaluation_function/models/__init__.py

@@ -12,6 +12,7 @@
     "shannon_words_ngram",
     "bengio_infer",
     "gpt2_infer",
+    "smollm2_infer",
 ]
 
 __all__ = ["AVAILABLE_MODELS"]

evaluation_function/models/gpt2_infer.py

@@ -3,127 +3,51 @@
 
 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)
 - ``temperature``: sampling temperature (default 1.0)
 - ``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)