Add GPT-2 model (gpt2_infer)

Dockerfile

@@ -30,6 +30,15 @@ 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.
+ENV GPT2_MODEL_DIR=/app/hf_models/gpt2 \
+    HF_HOME=/tmp/hf
+COPY scripts/download_gpt2.py /tmp/download_gpt2.py
+RUN HF_HOME=/tmp/hf_build python /tmp/download_gpt2.py "${GPT2_MODEL_DIR}" && \
+    rm -rf /tmp/hf_build /tmp/download_gpt2.py
+
 # Copy evaluation function first (smaller, changes more often)
 COPY evaluation_function ./evaluation_function
 

docs/dev.md

@@ -1,11 +1,21 @@
 # langModels
 An evaluation function to aid the learning of language models.
 
-**NOTE: **  All AWS deployments should allocate 1024MB RAM and Storage due to the image size.
+**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.
 
 ## Inputs
 *Specific input parameters which can be supplied when the `eval` command is supplied to this function.*
 
+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).
+
 ## Outputs
 *Output schema/values for this function*
 

docs/user.md

@@ -2,4 +2,16 @@
 
 A series of language models that students can run for inference to explore their performance.
 
-Created from scratch with just a neural network for modeling data from a sine wave. This is to ensure the software all works in integration tests. Following that, more models will be added (Shannon ngrams (letters); Shannon ngrams (words); Bengio neural network model of language with a small context window; basic transformer models; larger transformer models; external LLMs).
+Created from scratch with just a neural network for modeling data from a sine wave. This is to ensure the software all works in integration tests. Following that, more models will be added (Shannon ngrams (letters); Shannon ngrams (words); Bengio neural network model of language with a small context window; basic transformer models; larger transformer models; external LLMs).
+
+## Models
+
+### `gpt2_infer`
+Pre-trained GPT-2 (124M) from Hugging Face. The student's response is used as the prompt, and GPT-2's sampled continuation is returned as feedback.
+
+| Param | Default | Notes |
+|---|---|---|
+| `word_count` | 20 | New tokens to generate (1–100) |
+| `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 |

evaluation_function/dev.json

@@ -50,5 +50,13 @@
     "word_count": 12,
     "context_window": 3,
     "dev": true
+  },
+  "gpt2_infer": {
+    "response": "the general sweetness",
+    "answer": 0,
+    "model": "gpt2_infer",
+    "word_count": 20,
+    "temperature": 1.0,
+    "top_k": 50
   }
 }

evaluation_function/evaluation_test.py

@@ -1,4 +1,6 @@
+import os
 import unittest
+from unittest import mock
 
 from .evaluation import Params, evaluation_function
 
@@ -35,4 +37,29 @@ def test_evaluation(self):
         if isinstance(feedback, list) and feedback:
             self.assertIsInstance(feedback[0][1], str)
         else:
-            self.assertIsInstance(feedback, str)
+            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."""
+
+    @unittest.skipUnless(os.environ.get("RUN_GPT2_TESTS"), "set RUN_GPT2_TESTS=1 to run")
+    def test_generates_continuation(self):
+        prompt = "Once upon a time"
+        result = evaluation_function(
+            prompt, None, {"model": "gpt2_infer", "word_count": 5, "seed": 0}
+        ).to_dict()
+        self.assertTrue(result.get("is_correct"))
+        self.assertTrue(result["feedback"].startswith(prompt))
+        self.assertGreater(len(result["feedback"]), len(prompt))
+
+    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")))

evaluation_function/models/__init__.py

@@ -11,6 +11,7 @@
     "shannon_letters_ngram",
     "shannon_words_ngram",
     "bengio_infer",
+    "gpt2_infer",
 ]
 
 __all__ = ["AVAILABLE_MODELS"]

evaluation_function/models/gpt2_infer.py

@@ -0,0 +1,130 @@
+"""
+GPT-2 (124M) inference using the pre-trained Hugging Face ``gpt2`` checkpoint.
+
+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.
+
+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).
+
+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")
+
+DEFAULT_MODEL_ID = "gpt2"
+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))
+
+
+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)
+
+
+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)),
+            ],
+        )
+
+
+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"),
+    )
+
+    result = Result(True)
+    result.add_feedback("general", text)
+    return result