IvanHU/esmc-ar-846m-ple-wide16-no-gr-bidir-fim-5b

ESMC-AR 846M PLE wide16 bidirectional + FIM (5B tokens)

This is a fresh, completed 5B-token ESMC-AR variant with 16 wide Transformer blocks, PLE, gated attention and no gated residual. The FP32 safetensors export contains 845,824,000 parameters and is self-contained for transformers with bundled custom code and the ESMC tokenizer.

• —Training task: original next-token, reversed-row next-token, and PSM-formatted FIM tasks sampled approximately 1:1:1 by output tokens
• —Training data: fixed shuffled selection from QingWY/protein-pretraining-data, revision 23949d7e372cf5d0f002bb35dbb6f56d0fd63cb9
• —Final training token count: 5,000,000,000
• —Last reported original-stream validation loss: 2.358863 (task mixtures make cross-model comparisons non-equivalent)
• —Training source commit: 87d23e04cc64126e109c5b349c70fb54db9aa635
• —W&B run: https://wandb.ai/huyiwen/bio-esmc-ar/runs/a4aa0a215517
• —Context length: 2,048; vocabulary: ESMC token IDs; no automatic BOS/EOS insertion
• —FP32 export verified against the original training checkpoint; see verification.json

Usage

ESMC-AR Hugging Face inference

This directory is self-contained: config.json, configuration_esmc_ar.py, modeling_esmc_ar.py, feature/PLE Python dependencies, safetensors weight shards, tokenizer files, generation_config.json and export_info.json. Tested with PyTorch 2.5.1 and Transformers 4.51.0. Install torch, transformers, safetensors, tokenizers; FlashAttention is optional and the training repository is not required. Attention uses PyTorch SDPA.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

path = '/path/to/export'
tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModelForCausalLM.from_pretrained(
    path, trust_remote_code=True, torch_dtype=torch.bfloat16,
).to('cuda').eval()
batch = tokenizer(['MALWMRLLPLL', 'MKWVTFISLLFLFSSAYS'],
                  padding=True, return_tensors='pt').to(model.device)
with torch.inference_mode():
    ids = model.generate(**batch, max_new_tokens=64, do_sample=True,
                         temperature=0.8, top_p=0.95)
print(tokenizer.batch_decode(ids, skip_special_tokens=True))

trust_remote_code=True loads the bundled Python inference files. No network access is needed when the directory is local. For CPU, use torch_dtype=torch.float32 and omit .to('cuda'). FP32 weights under BF16 autocast most closely match the training numerics; directly casting all weights/residuals to BF16 trades a small numerical difference for inference speed and memory.

Supports left-padded batch generation, sampling/greedy/beam search, hidden states/attention outputs and standard shifted causal-LM labels. Baseline models also support inputs_embeds and a tuple KV cache with 2 KV heads. PLE exports require input_ids and use_cache=False; short-convolution exports also disable caching. For those variants, generation recomputes the full context. No packed training interface is exposed here: each batch row is one causal sequence.

Tokenizer IDs are ESMC IDs, not raw ProGen3 IDs. Automatic BOS/EOS insertion is disabled to match training. Supply a nonempty amino-acid prefix. The dataset consists of protein fragments and does not train biological end-of-protein termination. Generation therefore defaults to a bounded length and suppresses nonstandard amino acids, special tokens and unused vocabulary slots. To change this explicitly, pass suppress_tokens=None; this does not make previously unseen tokens trained.

RoPE was trained on fragments up to 2048 positions. Longer inference is mechanically possible, but its quality has not been evaluated. export_info.json records the source step/token count; an intermediate export is not a completed 5B-token model.

From the training repository, export any completed checkpoint with:

.venv/bin/python -m scripts.export_hf \
  --checkpoint /media/huyiwen/bio-esmc-ar-runs/stage1/checkpoint-00001250 \
  --output-dir /media/huyiwen/bio-esmc-ar-runs/hf-checkpoint-00001250