PubMed Pi 124M

Completed English biomedical base language model trained from scratch on PubMed titles and abstracts. It predicts and continues text; it is not instruction-tuned. No pretrained base model is used. Release v1.0.2, completed 2026-10-29 at step 82,452 after one epoch and 4,476,671,946 loss targets.

Use

Download this repository with huggingface_hub.snapshot_download("sraivante/pubmed-pi-124m"), install requirements.txt, then run from the downloaded directory:

python src/infer.py --model model.pt --tokenizer tokenizer.json --prompt "The role of inflammation in" --max-new-tokens 100

Python loading (run with the downloaded directory as working directory):

import sys, torch
from tokenizers import Tokenizer
sys.path.insert(0, "src")
from model import PubMedLM, ModelConfig
checkpoint = torch.load("model.pt", map_location="cpu", weights_only=True)
model = PubMedLM(ModelConfig(**checkpoint["config"]))
model.load_state_dict(checkpoint["model"])
model.eval()
tokenizer = Tokenizer.from_file("tokenizer.json")
tokenizer.encode_special_tokens = True
ids = [tokenizer.token_to_id("<bos>")] + tokenizer.encode("The role of inflammation in", add_special_tokens=False).ids
with torch.no_grad():
    output = model.generate(torch.tensor([ids]), max_new_tokens=80, temperature=0)
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))

This is a custom PyTorch architecture, not a Transformers AutoModel checkpoint. The supplied loading code is required. CPU and CUDA are supported; the evaluator was also run with PyTorch/XLA on TPU. The Hub inference widget is not configured.

Architecture and training

124,278,528 parameters; 14 decoder layers; width 768; 12 attention heads; MLP 3,072; RoPE; RMSNorm; GELU; tied embeddings; no biases; context 1,024; byte-level BPE vocabulary 32,768. Experimental initialization maps consecutive pi digits, including the initial 3, to float32 weights as 0.003 * digit. This is not a claim of superiority over standard initialization. Exact initialization manifests are included.

AdamW, betas (0.9, 0.95), epsilon 1e-8, weight decay 0.1, peak LR 3e-4, minimum LR 3e-5, 500 warmup steps, cosine schedule, microbatch 8, gradient accumulation 8, gradient clip 1.0, BF16 A100 training. Causal attention is isolated by document; padding and cross-document transitions do not contribute loss. Full settings, source, hashes, environment, training metrics and completion status are in provenance/ and src/.

Data and reproducibility

Versioned dataset companion contains split membership identifiers, exact hashes, filtering information and evaluation references. Training: 14,819,763 abstracts; validation: 74,281. The tokenizer used a deterministic train-only sample of 200,267 records. No fine-tuning dataset exists for this release.

The archive's training manifests identify the raw and packed files actually used. Historical upstream PubMed file revisions were not retained; current NLM downloads may differ. The companion distinguishes verified local snapshot identity from unknown upstream revision. Third-party abstract text and reversible token binaries are excluded because record-level redistribution rights have not been established. Open weights and source code do not imply that all training text is openly licensed.

Evaluation

HellaSwag evaluation used all 10,042 validation examples, four conditional ending likelihoods per row, explicit BOS, no EOS, FP32 on Colab v6e-1 TPU, PyTorch 2.8.0/PyTorch-XLA. Scoring took 130 seconds. Normalization divides log likelihood by ending character count, excluding the leading separator. Preprocessing follows EleutherAI's HellaSwag task; this is a custom adapter, not an official lm-evaluation-harness run. Standard errors: raw 0.004445, normalized 0.004504. Chance is 25%; these scores indicate limited general commonsense performance. Evaluation was performed after training, without fine-tuning on HellaSwag. No formal contamination audit was performed.

The exact result and evaluator are in evaluation/. Reproduce with the companion's evaluation snapshot (or the upstream source after checking its SHA-256):

python evaluation/hellaswag.py --model-dir . --data /path/to/hellaswag_val.jsonl --output results --device cpu

For TPU use matching torch==2.8.0 and torch_xla[tpu]==2.8.0, set PJRT_DEVICE=TPU and TORCH_DEVICE_BACKEND_AUTOLOAD=0, and pass --device xla.

Limitations and license

The model can hallucinate, repeat, and reflect outdated or biased abstracts. English is the intended language; other languages are unsupported. No labeled biomedical QA benchmark, clinical validation, instruction tuning, or safety alignment was performed. It is unsuitable for autonomous clinical decisions. Perplexity is not evidence of medical correctness.

Original model weights, code and documentation are released under Apache-2.0 by sraivante, to the extent the contributor holds rights. Third-party rights are preserved. NLM does not endorse this release. See LICENSE, NOTICE, and the companion dataset card for exclusions and NLM terms.

Download statistics

The root config.json contains the custom PubMedLM architecture settings and enables Hugging Face download counting. Download the complete repository using snapshot_download as shown above so this file is included. Requests for the configuration are counted server-side by Hugging Face; downloading only model.pt may not be counted. This configuration does not make the custom model Transformers AutoModel-compatible. See Hub download statistics.

Exact context and training rules

See TRAINING_RULES.md for the verified length filter, tokenizer, BOS/EOS, shifted targets, document-isolated packing, padding masks, long-document windowing, exact counts and optimization rules. A 1,024-position model context does not mean every abstract was exactly 1,024 tokens or that all long abstracts fit in one window.