Skip to content

Add GPT-2 model (gpt2_infer) - #6

Merged
m-messer merged 3 commits into
mainfrom
worktree-add-gpt2
Nov 4, 2026
Merged

m-messer merged 3 commits into
mainfrom
worktree-add-gpt2

Conversation

@m-messer

@m-messer m-messer commented Nov 3, 2026

Copy link
Copy Markdown
Member

Summary

Adds gpt2_infer, the next model after bengio_infer: pre-trained GPT-2 (124M) loaded with Hugging Face transformers. The student's response is used as the prompt, and GPT-2's sampled continuation is returned as general feedback.

  • Params: word_count (new tokens, default 20, max 100), temperature (default 1.0), top_k (default 50), seed (optional).
  • Weights: scripts/download_gpt2.py downloads them into /app/hf_models/gpt2 (GPT2_MODEL_DIR) at Docker build time, so nothing is fetched at runtime. Locally, they download on first use into HF_HOME (default /tmp/hf).
  • Failure handling: the same guard as bengio_infer. A load or inference error returns an error Result rather than crashing the RPC worker. The loaded model is cached, so a warm worker loads it once.
  • Dependency: adds transformers ^4.57 (lock file updated).

⚠️ Lambda memory

In the built image, with --memory 1024m, peak memory for one GPT-2 request was 989MB in the Python process alone. That's too close for a 1024MB Lambda once shimmy is counted. Recommend at least 2048MB before this model is used in production. docs/dev.md has been updated to say so.

Image size: 3.43 GB.

Testing

  • pytest evaluation_function/evaluation_test.py: 3 passed (with RUN_GPT2_TESTS=1, including a real download and generation).
    • The generation test is skipped unless RUN_GPT2_TESTS=1 is set, so CI doesn't download about 500MB.
    • A second test checks that a load failure returns an error result instead of raising.
  • Built with Podman, then ran in the container with --network none --memory 1024m. Weights loaded from /app/hf_models/gpt2 and generation worked.
  • python -m evaluation_function.dev --config dev.json --case gpt2_infer works when pointed at a local copy of the downloaded weights.

Note (not changed here)

evaluation_function/dev.py calls evaluation_function(answer, response, ...), swapping the two arguments. This affects every dev case. For gpt2_infer it means the dev case's answer (0) is the prompt, so it falls back to the default prompt.

🤖 Generated with Claude Code

m-messer and others added 2 commits November 3, 2026 17:47
Adds gpt2_infer, which uses the student's response as a prompt and
returns GPT-2's sampled continuation. Weights are baked into the Docker
image at build time so no network is needed at runtime; locally they
download on first use into HF_HOME.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
scripts/download_gpt2.py takes the target dir (default $GPT2_MODEL_DIR)
and an optional --model-id, so it can also be run locally.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

@peterbjohnson peterbjohnson left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great thanks. Minor comments.

Comment thread docs/dev.md Resolved
Comment thread docs/dev.md Resolved
@m-messer
m-messer merged commit a92dd24 into main Nov 4, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants