Match genome-scale inference windowing and batching

INFERENCE.md

@@ -18,11 +18,13 @@ python infer_genbank.py genome.fna.gz --format fasta --out-dir genome-prediction
 
 Both commands process every sequence record in the file. Use `--revision <model-commit>` to pin the model, or `--threshold 0.7` to change the binary-call threshold (default 0.5).
 
-The script uses 98,304-base windows with a 49,152-base stride and averages CDS probabilities where windows overlap. The last window may be shorter. No BOS/EOS tokens are added. Short records are supported, but the training data used full-length 98,304-base windows. This example uses overlap averaging; it does not reproduce any separate production genome-annotation pipeline.
+The script uses the same windowing and batching as Carbon-A's genome-scale predictions. Windows are 98,304 bases (16,384 tokens) and consecutive windows overlap by 6,144 bases (`--overlap-length`, default 1,024 tokens). The last window of a longer record is aligned to the record's end, so all windows are full length. A record shorter than one window is covered by one window from its start, of whole 6-mers, plus, when its length is not a multiple of six, a second one aligned to its end. CDS probabilities from the two strand heads are averaged where windows overlap. No BOS/EOS tokens or `N` bases are added; windows shorter than 16,384 tokens are padded with the pad token and masked. Records shorter than six bases get probability 0. Short records are supported, but the training data used full-length 98,304-base windows.
 
-Outputs are one compressed NumPy file per record, named with its accession and record name, plus `summary.json`. Each file contains forward/reverse CDS probabilities and calls, a combined `cds_label_any` track, a `valid` mask, and the number of contributing windows per base (`window_count`). Array position `i` corresponds to record base `i + 1` on **both** strands; the reverse track is not reversed. These are coding-occupancy predictions, not GFF gene or transcript models.
+In BF16, probabilities depend slightly on the attention kernel and on how windows are batched. The script therefore batches the same way: records are processed longest first in groups of four (`--dispatch-chunk-size`), and full-length windows and shorter, padded windows are run in separate batches of eight (`--batch-size`). The genome-scale predictions used FlashAttention-2. With `flash-attn`, Transformers 4.x, and `--attn-implementation flash_attention_2`, the script reproduced those probabilities bit for bit in our tests. The default, PyTorch SDPA, needs no extra package; in the same tests, single-base probabilities differed by up to 0.11, and at least 99.8% of 0.5 calls were unchanged. A smaller `--batch-size` uses less GPU memory but also changes results slightly.
 
-Any 6-mer containing an ambiguous nucleotide is excluded from the output. An incomplete final 6-mer is padded with `N` and also excluded. Unscored positions have probability `NaN`, call `-1`, and `valid=False`. Record sequences must be present in the file; annotation-only records with undefined sequence cannot be used. The example retains per-base arrays for one record in memory, so large chromosomes require more host RAM.
+Outputs are one compressed NumPy file per record, named with its accession and record name, plus `summary.json`. Each file contains forward/reverse CDS probabilities and calls, a combined `cds_label_any` track, an `ambiguous` mask for bases other than A, C, G, and T, and the number of contributing windows per base (`window_count`). Array position `i` corresponds to record base `i + 1` on **both** strands; the reverse track is not reversed. These are coding-occupancy predictions, not GFF gene or transcript models.
+
+Every base receives a probability. A 6-mer containing an ambiguous nucleotide is tokenized as an unknown token and still predicted; its predictions are less reliable. Record sequences must be present in the file; annotation-only records with undefined sequence cannot be used. The example retains per-base arrays for one record in memory, so large chromosomes require more host RAM.
 
 ```python
 import json
@@ -34,9 +36,9 @@ record = json.loads((out / "summary.json").read_text())["records"][0]
 with np.load(out / record["file"]) as prediction:
     plus = prediction["cds_probability_plus"]
     minus = prediction["cds_probability_minus"]
-    coding = prediction["cds_label_any"]  # 0, 1, or -1 (unscored)
-    valid = prediction["valid"]
-    print(record["accession"], plus.shape, int(valid.sum()))
+    coding = prediction["cds_label_any"]  # 1 if either strand passes the threshold
+    ambiguous = prediction["ambiguous"]
+    print(record["accession"], plus.shape, int(ambiguous.sum()))
 ```
 
 ## Run inference on packed data

README.md

@@ -48,9 +48,12 @@ Full-length 98,304-base windows rely on PyTorch's fused SDPA kernels (FlashAtten
 
 ### Use from Python
 
+Run from the downloaded repository, so that `infer_genbank.py` can be imported:
+
 ```python
 import torch
 from transformers import AutoModelForTokenClassification, AutoTokenizer
+from infer_genbank import predict_record
 
 model_id = "HuggingFaceBio/Carbon-A-1.2B"
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
@@ -58,16 +61,12 @@ model = AutoModelForTokenClassification.from_pretrained(
     model_id, trust_remote_code=True, dtype=torch.bfloat16, attn_implementation="sdpa"
 ).to("cuda").eval()
 
-sequence = "ATGTCTGCTACACAAAAGCAACACTAA"  # up to 98,304 bases
-sequence = sequence.upper()
-sequence += "N" * (-len(sequence) % 6)  # complete the final 6-mer
-inputs = tokenizer(sequence, add_special_tokens=False, return_tensors="pt").to("cuda")
-with torch.inference_mode():
-    logits = model(**inputs).logits  # (1, 2 * len(sequence), 2)
-cds_plus, cds_minus = logits.float().softmax(-1)[0, :, 1].reshape(2, -1).cpu()
+sequence = "ATGTCTGCTACACAAAAGCAACACTAA"  # any length
+probs, window_count, num_windows = predict_record(sequence, model, tokenizer, "cuda")
+cds_plus, cds_minus = probs  # per-base P(CDS) on the forward and reverse strand
 ```
 
-Logits are per base: all forward-strand positions, then all reverse-strand positions, both in input coordinates. Positions covered by a 6-mer containing `N` or another ambiguous base are not meaningful; the example scripts exclude them. Pass `attention_mask`, which the tokenizer returns, because the model requires it. Uppercase the sequence and use `add_special_tokens=False`: the tokenizer adds a BOS token by default, and lowercase (soft-masked) bases are not grouped into 6-mers, which shifts every later position.
+`predict_record` handles windowing, including records whose length is not a multiple of six. To call the model directly, pass one window: an uppercase sequence of whole 6-mers, at most 98,304 bases, tokenized with `add_special_tokens=False`, together with the `attention_mask` the tokenizer returns. The tokenizer otherwise adds a BOS token, and lowercase (soft-masked) bases are not grouped into 6-mers, which shifts every later position. The logits are per base: all forward-strand positions, then all reverse-strand positions, both in input coordinates.
 
 ### Annotate a GenBank or FASTA file
 
@@ -81,7 +80,7 @@ python infer_genbank.py genome.gbff.gz --out-dir predictions
 python infer_genbank.py genome.fna.gz --format fasta --out-dir predictions-fasta
 ```
 
-Long records are processed in **98,304-base windows with 50% overlap**. Probabilities are averaged where windows overlap. Existing GenBank annotations are not model inputs.
+Records are split into **98,304-base windows that overlap by 6,144 bases** (1,024 tokens). The last window is aligned to the end of the record, so every window is full length and no `N` bases are added. A record shorter than one window, whose length is not a multiple of six, is run twice: once aligned to its start and once to its end. Probabilities are averaged where windows overlap. Windowing and batching match Carbon-A's genome-scale predictions; see the inference guide for reproducing them exactly. Existing GenBank annotations are not model inputs.
 
 Each record produces an `.npz` file named with its accession and record name, plus an entry in `summary.json`. To read the predictions:
 
@@ -96,10 +95,10 @@ with np.load(out / record["file"]) as pred:
     plus = pred["cds_probability_plus"]
     minus = pred["cds_probability_minus"]
     coding = pred["cds_label_any"]  # 1 if either strand passes the threshold
-    valid = pred["valid"]
+    ambiguous = pred["ambiguous"]  # bases other than A, C, G, T
 ```
 
-Array index `i` corresponds to record base `i + 1` on both strands. Unscored positions have probability `NaN`, call `-1`, and `valid=False`.
+Array index `i` corresponds to record base `i + 1` on both strands. Every base has a probability; records shorter than six bases get 0.
 
 ### Run on packed training or evaluation data
 
@@ -130,7 +129,7 @@ The **[packed-data bucket](https://huggingface.co/buckets/HuggingFaceBio/Carbon-
 The inference examples have passed functional checks on an **NVIDIA H100**, using BF16, PyTorch 2.7.1, and Transformers 4.56.0:
 
 - Two full-length packed evaluation windows, including strand outputs and thresholded calls.
-- The complete **230,218-base yeast chromosome I** record, processed in four overlapping windows.
+- The complete **230,218-base yeast chromosome I** record.
 - Identical predictions when the same sequence is supplied as GenBank or FASTA.
 - Independent checks of overlap averaging, coordinates, ambiguous bases, and short records.
 
@@ -140,7 +139,7 @@ These checks verify the inference code. **They do not measure agreement with ref
 
 - Predictions represent **CDS occupancy**. Transcript paths, alternative isoforms, and complete gene structures are not resolved.
 - The model is intended for **eukaryotic sequence**. Bacteria and viruses are outside its training scope.
-- The GenBank/FASTA example excludes any 6-mer containing an ambiguous nucleotide, including an incomplete final 6-mer padded with `N`.
+- A 6-mer containing an ambiguous nucleotide (`N` or another IUPAC code) is an unknown token. It still receives predictions, but they are less reliable; the `ambiguous` array marks these bases.
 - Shorter records can be processed, but training used full-length 98,304-base windows.
 - The supplied evaluation panel was used during training; it is not a newly held-out test set.
 - Full-length inference on CPU is slow. The file example holds per-base arrays for one record in memory, so large chromosomes need more host RAM.

infer_genbank.py

@@ -11,59 +11,96 @@ import torch
 from transformers import AutoModelForTokenClassification, AutoTokenizer
 
 MODEL_ID = 'HuggingFaceBio/Carbon-A-1.2B'
-WINDOW = 98304
-STRIDE = 49152
+CONTEXT_TOKENS = 16384
+OVERLAP_TOKENS = 1024
+BATCH_SIZE = 8
+DISPATCH_CHUNK_SIZE = 4
+K = 6
 
 
 def safe_name(value):
     return re.sub(r'[^A-Za-z0-9_.-]+', '_', value).strip('._')[:100] or 'record'
 
 
-def predict_record(sequence, model, tokenizer, device, threshold):
-    """Average probabilities across overlapping windows in record coordinates."""
-    sequence = sequence.upper()
-    length = len(sequence)
-    if not length:
-        raise ValueError('Cannot predict an empty sequence')
-    sums = np.zeros((2, length), dtype=np.float32)
-    counts = np.zeros(length, dtype=np.uint32)
-    windows = 0
+def windows_for_sequence(length, window=CONTEXT_TOKENS * K, overlap=OVERLAP_TOKENS * K):
+    """(start, length) windows, each a whole number of 6-mers; the last is right-aligned."""
+    if length < window:
+        windows = [(0, length - length % K)] if length >= K else []
+        if length % K and length >= K:
+            windows.append((length % K, length - length % K))
+        return windows
+    windows, start, remaining = [], 0, length
+    while True:
+        size = min(window, remaining)
+        windows.append((start, size))
+        if size == remaining:
+            return windows
+        if remaining - window + overlap < window:
+            start, remaining = length - window, window
+        else:
+            start, remaining = start + window - overlap, remaining - (window - overlap)
+
+
+def predict_records(sequences, model, tokenizer, device, overlap_tokens=OVERLAP_TOKENS,
+                    batch_size=BATCH_SIZE, dispatch_chunk_size=DISPATCH_CHUNK_SIZE):
+    """Yield (index, probabilities, window_count, num_windows) for each sequence.
+
+    Probabilities have shape (2, length): forward, reverse strand. Overlapping windows are
+    averaged; sequences shorter than six bases get probability 0. Sequences are processed
+    longest first, in groups of `dispatch_chunk_size`, with full-length and padded windows
+    batched separately. Batch composition changes BF16 results slightly, so keep the
+    defaults to reproduce Carbon-A's genome-scale predictions.
+    """
+    sequences = [sequence.upper() for sequence in sequences]
+    order = sorted(range(len(sequences)), key=lambda i: len(sequences[i]), reverse=True)
+    pad_id = tokenizer.pad_token_id
     with torch.inference_mode():
-        for start in range(0, length, STRIDE):
-            stop = min(start + WINDOW, length)
-            fragment = sequence[start:stop]
-            # Complete the final 6-mer with N; its positions will be unscored.
-            fragment += 'N' * (-len(fragment) % 6)
-            encoded = tokenizer(fragment, add_special_tokens=False, return_tensors='pt',
-                                return_attention_mask=True)
-            ids = encoded['input_ids']
-            n_tokens = len(fragment) // 6
-            if ids.shape != (1, n_tokens):
-                raise ValueError('Tokenizer did not produce one token per six bases')
-            # Match packed-data inference: OOV tokens remain visible context;
-            # only their output positions are excluded, not the attention mask.
-            mask = torch.ones_like(ids)
-            logits = model(input_ids=ids.to(device), attention_mask=mask.to(device)).logits.float()
-            if logits.shape != (1, 2 * n_tokens * 6, 2) or not torch.isfinite(logits).all():
-                raise ValueError('Unexpected shape or nonfinite model predictions')
-            probs = logits.softmax(-1)[0, :, 1].cpu().numpy().reshape(2, -1)
-            width = stop - start
-            valid = np.repeat(ids[0].numpy() >= 32, 6)[:width]
-            sums[:, start:stop] += np.where(valid[None, :], probs[:, :width], 0)
-            counts[start:stop] += valid.astype(np.uint32)
-            windows += 1
-            if stop == length:
-                break
-    valid = counts > 0
-    np.divide(sums, counts[None, :], out=sums, where=valid[None, :])
-    sums[:, ~valid] = np.nan
-    calls = (sums >= threshold).astype(np.int8)
-    calls[:, ~valid] = -1
-    any_call = np.any(calls == 1, axis=0).astype(np.int8)
-    any_call[~valid] = -1
-    return {'cds_probability_plus': sums[0], 'cds_probability_minus': sums[1],
-            'cds_label_plus': calls[0], 'cds_label_minus': calls[1],
-            'cds_label_any': any_call, 'valid': valid, 'window_count': counts}, windows
+        for chunk_start in range(0, len(order), dispatch_chunk_size):
+            chunk = order[chunk_start:chunk_start + dispatch_chunk_size]
+            sums = {i: torch.zeros((2, len(sequences[i])), dtype=torch.float32) for i in chunk}
+            counts = {i: torch.zeros(len(sequences[i]), dtype=torch.long) for i in chunk}
+            num_windows = {}
+            full, padded = [], []
+
+            def run_batch(batch):
+                ids = [tokens + [pad_id] * (CONTEXT_TOKENS - len(tokens)) for *_, tokens in batch]
+                mask = [[1] * len(tokens) + [0] * (CONTEXT_TOKENS - len(tokens)) for *_, tokens in batch]
+                logits = model(input_ids=torch.tensor(ids, device=device),
+                               attention_mask=torch.tensor(mask, device=device)).logits
+                probs = logits.softmax(dim=-1).cpu()
+                for row, (i, start, size, _) in enumerate(batch):
+                    sums[i][:, start:start + size] += probs[row, :2 * size].view(2, size, -1)[:, :, 1].float()
+                    counts[i][start:start + size] += 1
+
+            for i in chunk:
+                windows = windows_for_sequence(len(sequences[i]), overlap=overlap_tokens * K)
+                num_windows[i] = len(windows)
+                for start, size in windows:
+                    tokens = tokenizer(sequences[i][start:start + size], add_special_tokens=False)['input_ids']
+                    if len(tokens) * K != size:
+                        raise ValueError(f'Tokenizer produced {len(tokens)} tokens for {size} bases')
+                    pending = full if len(tokens) == CONTEXT_TOKENS else padded
+                    pending.append((i, start, size, tokens))
+                    if len(pending) >= batch_size:
+                        run_batch(pending[:batch_size])
+                        del pending[:batch_size]
+            for pending in (full, padded):
+                while pending:
+                    run_batch(pending[:batch_size])
+                    del pending[:batch_size]
+
+            for i in chunk:
+                if num_windows[i] and torch.any(counts[i] == 0):
+                    raise RuntimeError('Some bases were not covered by any window')
+                probs = (sums[i] / counts[i].clamp(min=1).float()).numpy()
+                yield i, probs, counts[i].numpy().astype(np.uint32), num_windows[i]
+
+
+def predict_record(sequence, model, tokenizer, device, overlap_tokens=OVERLAP_TOKENS):
+    """Return (probabilities, window_count, num_windows) for one sequence; see predict_records."""
+    _, probs, counts, num_windows = next(predict_records([sequence], model, tokenizer, device,
+                                                         overlap_tokens))
+    return probs, counts, num_windows
 
 
 def main():
@@ -74,10 +111,22 @@ def main():
     parser.add_argument('--revision', default='main', help='Use a model commit for repeatable runs.')
     parser.add_argument('--device', choices=['cuda', 'cpu'], default='cuda')
     parser.add_argument('--threshold', type=float, default=0.5)
+    parser.add_argument('--overlap-length', type=int, default=OVERLAP_TOKENS,
+                        help='Overlap between consecutive windows, in 6-mer tokens.')
+    parser.add_argument('--batch-size', type=int, default=BATCH_SIZE,
+                        help='Windows per forward pass; lower it if GPU memory is short.')
+    parser.add_argument('--dispatch-chunk-size', type=int, default=DISPATCH_CHUNK_SIZE,
+                        help='Records whose windows are batched together.')
+    parser.add_argument('--attn-implementation', choices=['sdpa', 'flash_attention_2'], default='sdpa',
+                        help='flash_attention_2 requires the flash-attn package.')
     parser.add_argument('--out-dir', type=Path, default=Path('genbank-predictions'))
     args = parser.parse_args()
     if not 0 <= args.threshold <= 1:
         parser.error('Threshold must be between 0 and 1.')
+    if not 0 <= args.overlap_length < CONTEXT_TOKENS:
+        parser.error(f'Overlap length must be between 0 and {CONTEXT_TOKENS - 1} tokens.')
+    if args.batch_size < 1 or args.dispatch_chunk_size < 1:
+        parser.error('Batch size and dispatch chunk size must be positive.')
     if not args.input.is_file():
         parser.error('Input file does not exist.')
     if args.device == 'cuda' and not torch.cuda.is_available():
@@ -87,35 +136,51 @@ def main():
     dtype = torch.float32
     if args.device == 'cuda':
         dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    opener = gzip.open if args.input.suffix.lower() == '.gz' else open
+    with opener(args.input, 'rt') as handle:
+        records = [(record.id, record.name, str(record.seq).upper())
+                   for record in SeqIO.parse(handle, args.format)]
+    if not records:
+        parser.error('No sequence records found; check the input format.')
+    for accession, _, sequence in records:
+        if not sequence:
+            raise ValueError(f'Empty sequence in record {accession}')
     tokenizer = AutoTokenizer.from_pretrained(args.model, revision=args.revision, trust_remote_code=True)
     model = AutoModelForTokenClassification.from_pretrained(
         args.model, revision=args.revision, trust_remote_code=True,
-        dtype=dtype, attn_implementation='sdpa',
+        dtype=dtype, attn_implementation=args.attn_implementation,
     ).to(args.device).eval()
-    if (model.config.k, model.config.num_prediction_heads, model.config.num_labels) != (6, 2, 2):
+    if (model.config.k, model.config.num_prediction_heads, model.config.num_labels) != (K, 2, 2):
         raise ValueError('Expected a 6-mer, two-strand binary CDS model.')
     args.out_dir.mkdir(parents=True, exist_ok=True)
     summary = {'model':args.model, 'revision_requested':args.revision,
                'revision_resolved':getattr(model.config, '_commit_hash', None),
                'input_file':str(args.input), 'input_format':args.format,
-               'dtype':str(dtype), 'threshold':args.threshold,
-               'window_bp':WINDOW, 'stride_bp':STRIDE,
+               'dtype':str(dtype), 'attn_implementation':args.attn_implementation,
+               'threshold':args.threshold, 'window_bp':CONTEXT_TOKENS * K,
+               'overlap_bp':args.overlap_length * K, 'batch_size':args.batch_size,
+               'dispatch_chunk_size':args.dispatch_chunk_size,
                'coordinates':'Array index i corresponds to record base i+1 on both strands.',
                'records':[]}
-    opener = gzip.open if args.input.suffix.lower() == '.gz' else open
-    with opener(args.input, 'rt') as handle:
-        for index, record in enumerate(SeqIO.parse(handle, args.format)):
-            if not len(record):
-                raise ValueError(f'Empty sequence in record {record.id}')
-            arrays, windows = predict_record(str(record.seq), model, tokenizer, args.device, args.threshold)
-            filename = f'{index:04d}_{safe_name(record.id)}_{safe_name(record.name)}.npz'
-            np.savez_compressed(args.out_dir / filename, **arrays)
-            summary['records'].append({'accession':record.id, 'name':record.name,
-                                       'length':len(record), 'windows':windows,
-                                       'valid_bases':int(arrays['valid'].sum()), 'file':filename})
-            print(f'Saved {filename}: {len(record):,} bases, {windows} windows', flush=True)
-    if not summary['records']:
-        raise ValueError('No sequence records found; check the input format.')
+    entries = {}
+    predictions = predict_records([sequence for *_, sequence in records], model, tokenizer,
+                                  args.device, args.overlap_length, args.batch_size,
+                                  args.dispatch_chunk_size)
+    for index, probs, counts, windows in predictions:
+        accession, name, sequence = records[index]
+        calls = (probs >= args.threshold).astype(np.int8)
+        ambiguous = ~np.isin(np.frombuffer(sequence.encode(), dtype=np.uint8),
+                             np.frombuffer(b'ACGT', dtype=np.uint8))
+        arrays = {'cds_probability_plus': probs[0], 'cds_probability_minus': probs[1],
+                  'cds_label_plus': calls[0], 'cds_label_minus': calls[1],
+                  'cds_label_any': calls.max(axis=0), 'ambiguous': ambiguous,
+                  'window_count': counts}
+        filename = f'{index:04d}_{safe_name(accession)}_{safe_name(name)}.npz'
+        np.savez_compressed(args.out_dir / filename, **arrays)
+        entries[index] = {'accession':accession, 'name':name, 'length':len(sequence),
+                          'windows':windows, 'ambiguous_bases':int(ambiguous.sum()), 'file':filename}
+        print(f'Saved {filename}: {len(sequence):,} bases, {windows} windows', flush=True)
+    summary['records'] = [entries[index] for index in sorted(entries)]
     (args.out_dir / 'summary.json').write_text(json.dumps(summary, indent=2)+'\n')