IvanHU/esmc-ar-stage1-ple-wide16-shuffled-bidir

ESMC-AR Hugging Face inference

This directory is self-contained: config.json, configuration_esmc_ar.py, modeling_esmc_ar.py, model.safetensors, 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 two Python files in this export. 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 KV cache, left-padded batch generation, sampling/greedy/beam search, inputs_embeds, hidden states/attention outputs and standard shifted causal-LM labels. A returned cache is a tuple of per-layer (K,V) tensors with 2 KV heads. 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