README.md

metadata

library_name: transformers
tags:
  - genomics
  - dna
  - gene-annotation
  - ntv3
  - long-range
license: other
language:
  - code

TEMPORARY FILE — DELETE BEFORE RELEASE. This README is a working record of the JAX → PyTorch port, written for the NTv3 release team while the repository is being assembled. It is not the model card. Replace it in full with the public model card before this repository is made public.

NTv3 Annotate — plants (staging)

PyTorch conversion of new_ntv3_700M_annotate_plants_1mb, the released gene-annotation model for this clade. 703,551,441 parameters across 511 tensors (477 nnx.Param plus 34 nnx.BatchStat). Modelling code is not in this repository: config.json resolves it from InstaDeepAI/NTv3_pretrained through auto_map, so loading requires trust_remote_code=True.

Source checkpoint: s3://genomix-research-3e09d/checkpoints/ntv3_annotate/new_ntv3_700M_annotate_plants_1mb (MLflow run 5f3370cb757341b3bcd2acb064c033c1, step 877404, ≈920 B tokens, trained on 53 plant assemblies).

Verification against the JAX reference

Measured on the 1024 bp fixture the checkpoint bundle ships in its own test_inputs item (sequence = "ACGT" * 256), comparing the annotation head's bed_tracks_logits of shape (1, 1024, 1, 15).

comparison	max abs	mean abs	decoded-class agreement
torch fp32 vs JAX fp32	3.386e-05	5.086e-06	100%
torch bf16 autocast vs JAX fp32	8.624e-02	1.914e-02	100%

Conclusion: the port is correct. The fp32 agreement is accumulation noise across a 703M-parameter network. The bf16 row is dominated by the intrinsic cost of bf16 rather than by the conversion: on the mammals model, where a true JAX bf16 reference was also generated, JAX bf16 differs from JAX fp32 by 6.28e-02 on its own, and torch bf16 differs from JAX bf16 by 6.82e-02 — the same magnitude, so the port adds no error of its own. Most importantly the decoded class is identical at every one of the 1024 positions in both precisions, and since a Viterbi decode consumes the argmax, neither setting can change the emitted GFF3.

The converter refuses to write weights unless decoded-class agreement is 100% in both precisions, so this file could not have been produced from a failing conversion.

Precision guidance

Weights ship in fp32. bf16 is a compute dtype, matching the JAX build_bf16_compute_model path, which lowers five trunk stages (stem, down_convolution, up_convolution, transformer_qkvo, transformer_ffn) and leaves parameters, norms, embeddings, the LM head and the annotation head in fp32:

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    out = model(tokens, conditions_tokens=conds)

Do not use torch_dtype="bfloat16" or model.to(torch.bfloat16). A whole-module cast raises expected m1 and m2 to have the same dtype: RMSBatchNorm.forward promotes its input to ≥fp32 and never casts back, StandardizedConv1D then follows the input dtype, and the first bare nn.Linear (GQABlock.q_layer) still holds bf16 weights.

How the weights were converted

Flax variables map one-to-one onto the torch state_dict, loaded with strict=True. Leaf renames are kernel/embedding/scale → weight, under a model. prefix. Layout rules are keyed on the owning torch module, never on the tensor shape: seven kernels (decoder_blocks.{0..6}.pointwise_unet_skip.proj.kernel) are square because each maps C → C, so a shape-based rule cannot tell a transpose from a copy. Copying them produces a model that loads cleanly under strict=True, carries the correct parameter count, and still agrees 100% on the annotation argmax while the trunk is silently wrong from the first decoder block. Census check: the model has 108 Linear.weight tensors; a converter that transposes 101 has this bug.

Known gaps in this staging repository

• Tokenizer. vocab.json here is the 16-token soft-masked vocabulary these models were trained on (A T C G N plus lowercase a t c g n), which differs from the 11-token vocabulary of the initial NTv3 release. tokenizer_config.json resolves the tokenizer class from InstaDeepAI/NTv3_pretrained through auto_map, so AutoTokenizer.from_pretrained(..., trust_remote_code=True) returns a 16-token tokenizer reading this repository's own vocab.json.
• No decoding code. config.json carries the full HMM specification under annotation — 15 states, the transition grammar, motif rules, the U12 lane, the served crop of 0.75 and the 1,048,576 bp window — but the Viterbi decoder and GFF3 writer are not published here yet. Logits alone are not a gene annotation.
• Inputs are soft-masked. Lowercase repeat-masked bases carry their own token ids; an uppercase-only assembly loses the repeat signal without raising.
• Window geometry is not recoverable from the weights. ntv3.keep_target_center_fraction is set to the validated serving crop of 0.75, so a default forward reproduces jax_model/roundtrip_fixture.npz exactly. Running at a different window produces a plausible but wrong annotation rather than an error. Note that annotation.window.trained_keep_target_center_fraction records 1.0 while jax_model/hyperparameters.json records 0.75 for the same field; that disagreement is unresolved and does not affect serving, since both permit 0.75.
• The model is unconditioned. It carries no condition table, so forward takes an empty conditions_tokens list, the same as the other two annotation repositories.
• transformers 5. residual_scale is stored as a scalar, as the model declares it; it was a one-element list before, which transformers 5 refused to load.