genatator-pipeline

GENATATOR-PIPELINE

GENATATOR-PIPELINE is a Hugging Face pipeline for ab initio gene annotation from genomic DNA. It accepts a FASTA file, finds candidate transcript intervals, assigns transcript type, predicts exon and CDS structure, and writes a GFF3 annotation file.

The pipeline combines interval discovery, transcript-type classification, segmentation, filtering, and GFF generation in one transformers.pipeline call. The output of the call is a Python string containing the path to the written GFF file.

GENATATOR-PIPELINE supports CUDA GPU execution only and float32 inference only. CPU execution of the stage models and lower precision modes are not supported. CPU RAM holds model caches and prediction buffers. Preprocessing, reconstruction, filtering, and GFF generation run on CPU.

Inference can use one GPU, multiple GPUs on one node, or multiple nodes in a Slurm allocation. GPU workers process chromosome-stage assignments, while CPU owners manage the annotation workflow for each chromosome. Completed milestones are saved for resumption, which is enabled by default.

Hugging Face pipeline usage

Basic example

from transformers import pipeline

if __name__ == "__main__":
    pipe = pipeline(
        task="genatator-pipeline",
        model="AIRI-Institute/genatator-pipeline",
        trust_remote_code=True,
        device=0,
        dtype="float32",
    )

    output_path = pipe(
        "genome.fasta",
        output_gff_path="genome.gff",
    )
    print(output_path)

Save Python launch examples as scripts and keep the pipeline creation and call inside the if __name__ == "__main__": guard. GPU workers use spawned processes. The basic example selects visible CUDA device 0. By default, resume=True and the four batch sizes are set to 32, while RMT segmentation remains limited to 1.

Multiple GPUs on one node

from transformers import pipeline

if __name__ == "__main__":
    pipe = pipeline(
        task="genatator-pipeline",
        model="AIRI-Institute/genatator-pipeline",
        trust_remote_code=True,
        dtype="float32",
        devices=[0, 1, 2, 3],
    )

    output_path = pipe(
        "genome.fasta",
        output_gff_path="genome.gff",
        checkpoint_dir="genome-checkpoints",
        cpu_ram_limit_gib_per_node=100,
    )
    print(output_path)
    print(pipe.last_summary["complete"])
    print(pipe.last_summary["totals"])

Set the RAM budget below the memory available to the job. For example, the command above uses 100 GiB. Omit the argument for automatic detection. An explicit devices list takes precedence over device. Device indices refer to the GPUs visible to the process through CUDA_VISIBLE_DEVICES.

For a local repository, use its directory as model, for example model="/shared/genatator-pipeline". Multi-node execution uses the Slurm command-line launch.

Example with all main supported parameters defined

from transformers import pipeline

if __name__ == "__main__":
    pipe = pipeline(
        task="genatator-pipeline",
        model="AIRI-Institute/genatator-pipeline",
        trust_remote_code=True,
        devices=[0, 1, 2, 3],
        parallel=True,
        dtype="float32",
        edge_model_path="AIRI-Institute/genatator-moderngena-base-multispecies-edge-model",
        region_model_path="AIRI-Institute/genatator-moderngena-base-multispecies-region-model",
        transcript_type_model_path="AIRI-Institute/genatator-caduceus-ps-multispecies-transcript-type",
        segmentation_model_path="AIRI-Institute/genatator-caduceus-ps-multispecies-segmentation",
        edge_context_length=1024,
        region_context_length=8192,
        transcript_type_context_length=250000,
        segmentation_context_length=250000,
        edge_average_token_length=9.0,
        region_average_token_length=9.0,
        edge_max_genomic_chunk_ratio=1.5,
        region_max_genomic_chunk_ratio=1.5,
        edge_drop_last=False,
        region_drop_last=False,
        edge_apply_sigmoid=False,
        region_apply_sigmoid=False,
        transcript_type_apply_sigmoid=True,
        segmentation_apply_sigmoid=True,
        edge_gap_token_id=5,
        region_gap_token_id=5,
    )

    output_path = pipe(
        "genome.fasta",
        output_gff_path="genome.gff",
        edge_batch_size=32,
        region_batch_size=32,
        transcript_type_batch_size=32,
        segmentation_batch_size=32,
        segmentation_is_rmt=None,
        resume=True,
        checkpoint_dir="genome-checkpoints",
        cpu_ram_limit_gib_per_node=None,
        disk_offload_limit_gib_total=None,
        offload_dir=None,
        result_buffer_gib=1.0,
        max_active_chromosomes_per_node=None,
        edge_prefetch_chunks=1,
        seed=42,
        max_task_retries=1,
        worker_timeout_seconds=300.0,
        startup_timeout_seconds=900.0,
        cpu_threads_per_worker=1,
        cpu_batch_workers_per_gpu=4,
        cpu_prefetch_batches=2,
        cpu_prefetch_buffer_gib_per_gpu=0.25,
        model_revisions=None,
        local_files_only=False,
        reference_gff_path=None,
        summary_path=None,
        edge_context_fraction=0.5,
        region_context_fraction=0.5,
        gene_finding_use_reverse_complement=True,
        transcript_type_use_reverse_complement=True,
        segmentation_use_reverse_complement=True,
        lp_frac=0.05,
        pk_prom=0.1,
        pk_dist=50,
        pk_height=None,
        interval_window_size=2_000_000,
        max_pairs_per_seed=10,
        gene_finding_global_chunk_size=70_000_000,
        prob_threshold=0.5,
        zero_fraction_drop_threshold=0.01,
        transcript_type_threshold=0.5,
        splice_filter=True,
        deduplicate=True,
        intronic_filtering=True,
        keep_longest_terminal_variant=True,
        predict_internal_structure=True,
        transcript_coloring_thresholds="auto",
        use_cds_heuristic=True,
        save_intermediate_files=False,
        intermediate_output_dir=None,
        pairing_progress_every=1000,
        chunk_log_every=1000,
        shift=None,
    )
    print(output_path)

Each batch-size argument is a per-GPU maximum, not a chromosome count. Edge and region batches contain model-input windows. Transcript-type batches contain interval inputs. Segmentation distributes whole intervals to workers and batches their internal inference windows locally. RMT segmentation has an effective maximum of 1. Exact-length bucketing, partial batches, and automatic out-of-memory recovery can also produce smaller batches. Use the stage-specific batch-size arguments above, rather than the generic Transformers batch_size argument, to control model-window batching.

Parameter reference

Model paths, dtype, and parallel are selected when creating the Hugging Face pipeline. Biological processing parameters and parallel-runtime settings can be stored at construction or supplied in the pipeline call. Call-time values override the stored defaults for that call. The parameter descriptions below apply to the default parallel=True runtime unless stated otherwise.

For the command line, biological parameters and model paths belong in --config. Runtime settings can also be supplied in --runtime-config or through their hyphenated flags, for example --edge-batch-size. Command-line flags override JSON settings.

Pipeline setup arguments

• task — Hugging Face task name for this custom pipeline. Use "genatator-pipeline".
• model — Hugging Face repository or local directory containing the GENATATOR pipeline wrapper. For the published version, use "AIRI-Institute/genatator-pipeline".
• trust_remote_code — Must be True because the pipeline uses custom Python code from the model repository. Without it, Transformers will not load the custom pipeline class.
• device — Single CUDA GPU selection, for example 0 for the first visible GPU. Use devices to select multiple GPUs. In parallel mode the lightweight Hugging Face wrapper remains on CPU, but all stage-model forward passes execute on CUDA. Therefore, device=-1 is not a CPU-inference mode.
• dtype — Tensor dtype used when loading the stage models. Only "float32" is supported.

Model repositories

• edge_model_path — Repository or local path for the edge model. This model predicts transcript boundary signals, namely TSS and PolyA signals on both strands.
• region_model_path — Repository or local path for the region model. This model predicts strand-specific intragenic signal that is used to remove weak candidate transcript intervals.
• transcript_type_model_path — Repository or local path for the transcript-type classifier. This model labels each retained candidate interval as mRNA or lnc_RNA.
• segmentation_model_path — Repository or local path for the segmentation model. This model predicts the internal exon, intron, and CDS structure of each retained interval.

Context-length and chunking parameters

• edge_context_length — Token length of each edge-model input window, including tokenizer system tokens. Larger values give the edge model more context, but also increase memory use.
• region_context_length — Token length of each region-model input window, including tokenizer system tokens. The default is 8192 tokens, matching the gene-finding benchmark and manuscript configuration.
• transcript_type_context_length — Maximum token length passed to the transcript-type model for each candidate interval and evaluated orientation. Only the leading prefix up to this context is classified. With reverse-complement inference enabled, the reverse complement is formed from the whole interval before that orientation's prefix is selected.
• segmentation_context_length — Nucleotide length of each segmentation-model block inside a retained interval. The planner uses zero regular overlap. If needed, it shifts the final block back to cover the interval end. Predictions at positions covered by more than one block are averaged.
• edge_average_token_length — Estimated average number of nucleotides represented by one edge-model tokenizer token. The pipeline uses this value to convert token context length into genomic window length before tokenization.
• region_average_token_length — Estimated average number of nucleotides represented by one region-model tokenizer token. The pipeline uses this value to convert token context length into genomic window length before tokenization.
• edge_max_genomic_chunk_ratio — Maximum expansion ratio for edge-model genomic extraction before tokenizer truncation. It gives the tokenizer extra nucleotide sequence so the final tokenized window can be filled reliably.
• region_max_genomic_chunk_ratio — Maximum expansion ratio for region-model genomic extraction before tokenizer truncation. It plays the same role as edge_max_genomic_chunk_ratio, but for the region model.
• edge_context_fraction — Fractional overlap between consecutive edge-model genomic windows. Higher overlap can smooth boundary predictions, but it increases the number of model calls.
• region_context_fraction — Fractional overlap between consecutive region-model genomic windows. Higher overlap can make intragenic masks more stable, but it increases computation time.
• edge_drop_last — If True, the final incomplete edge-model window is omitted. The default is False, which keeps the final window so the end of the sequence is still processed.
• region_drop_last — If True, the final incomplete region-model window is omitted. The default is False, which keeps the final window so the end of the sequence is still processed.
• edge_gap_token_id — Token ID used to correct edge-model offset mappings for gap tokens. Most users should keep the default unless they change the tokenizer.
• region_gap_token_id — Token ID used to correct region-model offset mappings for gap tokens. Most users should keep the default unless they change the tokenizer.
• gene_finding_global_chunk_size — Maximum nucleotide length of each global chunk used by the edge model only. CPU reconstruction, FFT smoothing, and peak calling complete inside each global chunk. Sparse peak coordinates are retained, and dense edge arrays are released. CPU processing advances in chunk order. GPU work may look ahead according to edge_prefetch_chunks. The default is 70,000,000 nucleotides.

Interval-discovery parameters

• lp_frac — Fraction of the Fourier spectrum retained by the low-pass smoother before peak detection. Smaller values produce smoother boundary tracks and can remove local noise.
• pk_prom — Minimum peak prominence used during TSS and PolyA boundary detection. Higher values make peak calling more conservative.
• pk_dist — Minimum nucleotide distance between neighboring peaks of the same boundary class. This helps avoid calling several nearby peaks for one broad signal.
• pk_height — Optional minimum peak height after smoothing. Use None to disable this extra height filter.
• interval_window_size — Maximum distance allowed when pairing a TSS peak with a PolyA peak on the same strand. Candidate transcript intervals longer than this pairing window are not created.
• max_pairs_per_seed — Maximum number of nearest PolyA partners retained for each TSS seed. Larger values create more candidate intervals and can increase downstream computation.
• prob_threshold — Threshold used to convert region-model intragenic signal into a binary mask. A base is considered intragenic only when the model signal is above this threshold.
• zero_fraction_drop_threshold — Maximum tolerated fraction of non-intragenic bases inside a candidate interval. Intervals with a larger fraction below prob_threshold are discarded.

Reverse-complement options

• gene_finding_use_reverse_complement — Enables reverse-complement averaging for the edge and region models. This can improve strand-aware interval discovery, but it roughly doubles gene-finding model calls.
• transcript_type_use_reverse_complement — Enables reverse-complement averaging for transcript-type classification. The forward and reverse-complement scores are averaged before the final mRNA or lnc_RNA decision.
• segmentation_use_reverse_complement — Enables reverse-complement averaging for segmentation. This can stabilize structure prediction, but it increases segmentation compute cost.

Activation options

• edge_apply_sigmoid — Applies an additional sigmoid to edge-model output channels before token-to-nucleotide projection. The default is False, because the published edge model outputs are already expected in the correct scale.
• region_apply_sigmoid — Applies an additional sigmoid to region-model output channels before token-to-nucleotide projection. The default is False, because the published region model outputs are already expected in the correct scale.
• transcript_type_apply_sigmoid — Applies sigmoid to single-logit transcript-type outputs before thresholding. For multi-logit outputs, the pipeline uses softmax instead.
• segmentation_apply_sigmoid — Applies an additional sigmoid to segmentation-model output channels before structural decoding. The default is True for the published segmentation setup.

Transcript and segmentation parameters

• transcript_type_threshold — Threshold applied to the predicted lnc_RNA probability. Intervals at or above this value are labeled lnc_RNA, and intervals below it are labeled mRNA.
• splice_filter — Enables splice-motif filtering and terminal splice-boundary correction for exon and CDS segments. This post-processing step can remove or adjust segments that disagree with expected splice signals.
• deduplicate — Removes duplicate final transcript predictions. This is applied near the end of GFF generation to avoid repeated identical transcripts.
• intronic_filtering — Drops transcript predictions whose segmentation starts or ends with the intron class. This removes predictions that appear to begin or end inside an intron.
• keep_longest_terminal_variant — For overlapping transcripts with the same internal structure, keeps the longest terminal variant. This reduces redundant terminal variants that differ mainly by transcript ends.
• predict_internal_structure — Controls whether the pipeline continues past interval discovery into transcript-type classification, segmentation, and GFF generation. Keep it as True for normal annotation output. With False, filtered intervals are saved in checkpoints and optional intermediate files, no GFF is written, and the call returns an empty string.
• use_cds_heuristic — Uses an exon-derived CDS heuristic to select coding regions for mRNA transcripts. With False, CDS segments come from segmentation predictions. This affects mRNA transcripts only, and no CDS is emitted for lnc_RNA transcripts.
• transcript_coloring_thresholds — Controls transcript color bins in the output GFF. Use "auto" to split the observed segmentation-confidence range into four bins, or provide a custom list of exactly four thresholds.

The hardcoded transcript color map is applied to the final transcript set after structural filtering and optional CDS heuristic processing, followed by longest-terminal-variant selection and deduplication.

• Lowest bin, #66cc66, light green.
• Second bin, #006400, dark green.
• Third bin, #dcdcff, light blue.
• Top bin, #0c0c78, dark blue.

Intermediate-output, logging, and coordinate parameters

• save_intermediate_files — If True, writes gene-finding intermediate artifacts for each FASTA record: compact edge peak .npz files, intragenic-mask .npz files when region masks are available, .bed interval files, and an .h5 debug dump when h5py is installed. Filenames include the FASTA record ordinal and a sanitized record name. These files are separate from resume checkpoints.
• intermediate_output_dir — Output directory for intermediate artifacts. If omitted, intermediate files are written next to the input FASTA file.
• pairing_progress_every — Logging interval, measured in TSS seeds, during candidate interval construction. Increase it for less frequent logs.
• chunk_log_every — Logging interval used by the direct edge and region inference helpers. Parallel workers also emit stage, batch, and checkpoint events. This value does not set the frequency of those runtime events.
• shift — Coordinate offset applied to final GFF coordinates. Use an integer offset directly, or use "UCSC" to infer the offset from FASTA headers of the form chrom:start-end.
• output_gff_path — Path of the GFF file written by the pipeline call. If omitted, the input FASTA suffix is replaced with .gff. In parallel mode the file is published after finalization, rather than appended as chromosomes finish. Discovery-only runs do not write a GFF.

GPU selection and execution

• parallel — Enables GPU-worker execution. The default is True, including when only one GPU is selected. Set this at pipeline construction, not per call. False selects a single-GPU diagnostic mode with batch size 1. Worker scheduling, resumption, and parallel-runtime summaries do not apply to that mode.
• devices — List of distinct visible CUDA device indices, for example [0, 1, 2, 3]. The default is None: the runtime uses device when supplied through the Hugging Face interface, otherwise all visible GPUs. Under Slurm, the list is interpreted on each node. Every node must expose the requested local indices.
• edge_batch_size — Maximum number of edge-model input rows per GPU forward call. Default: 32.
• region_batch_size — Maximum number of region-model input rows per GPU forward call. Default: 32.
• transcript_type_batch_size — Maximum number of transcript-type input rows per GPU forward call. Default: 32. Forward and reverse-complement inputs are evaluated and combined for each interval when enabled.
• segmentation_batch_size — Maximum number of non-RMT segmentation input rows per GPU forward call. Default: 32. Whole intervals stay assigned to one worker. Their internal windows can be batched on that worker.
• segmentation_is_rmt — Optional explicit RMT restriction. None uses model, configuration, and module detection. True forces segmentation batch size 1. A detected RMT restriction remains effective even if this argument is False.
• seed — Seed for random equal-count assignment of transcript-type and segmentation intervals to workers. Default: 42. It controls work ownership, not the order of CPU annotation decisions.
• cpu_threads_per_worker — PyTorch CPU-operation thread limit within each GPU process. Default: 1. This is separate from the threads that prepare input batches and from chromosome-owner CPU processing.
• cpu_batch_workers_per_gpu — Requested maximum number of CPU input-preparation threads inside each GPU process. Default: 4. The runtime lowers this maximum when Linux CPU affinity, the Slurm task allocation, or a cgroup CPU quota leaves insufficient capacity. A value of 0 selects synchronous input preparation. The effective thread count is logged for each node and GPU.
• cpu_prefetch_batches — Maximum number of future batches whose input rows may be prepared in advance. Default: 2. A value of 0 disables lookahead but keeps parallel preparation within the current batch. The current chromosome-stage release gate and result-credit prefix can shorten the available lookahead.
• cpu_prefetch_buffer_gib_per_gpu — CPU input-cache allowance per GPU process, in GiB. Default: 0.25. Both an item-count limit and conservative byte reservations bound queued, preparing, and cached inputs. This allowance is included in the node's managed-memory reserve. It does not include active-batch tensors, tokenizer objects, or every temporary allocation made inside a tokenizer. An individual input that exceeds the allowance is prepared synchronously without shortening it.

Memory and offloading parameters

All GiB values use 1 GiB = 1024**3 bytes.

• cpu_ram_limit_gib_per_node — Per-node managed-memory admission budget. Default: None, which selects a conservative limit from available memory and detected Slurm/cgroup limits. The runtime reserves capacity for worker model caches and transport before admitting chromosome workspaces. This is an estimate rather than an operating-system-enforced RSS cap. Leave headroom below the job's hard memory limit.
• result_buffer_gib — Global budget for coordinator result credits. Default: 1.0 GiB. Credits are reserved before dispatch and returned as results are consumed. An individual result larger than this budget causes a reported stage failure. Increase the budget to accommodate it. This does not cap every CPU array, checkpoint, or chromosome workspace.
• max_active_chromosomes_per_node — Maximum number of chromosome CPU owners admitted on a node. Default: None, which uses that node's selected GPU count. Actual concurrency is also constrained by available GPUs, RAM, and active work.
• edge_prefetch_chunks — Number of additional global edge chunks whose GPU work may run ahead of the CPU consumer. Allowed values are 0 or 1. The default is 1. CPU reconstruction and peak calling still complete in chunk order.
• disk_offload_limit_gib_total — Run-wide temporary offload budget. Default: None, which disables temporary disk spilling. A positive integer permits up to that many GiB of run-owned temporary mapped arrays across nodes. Persistent milestone checkpoints are not included in this budget.
• offload_dir — Shared writable root for temporary offload files. Default: None, which places them under the run's checkpoint namespace when offloading is enabled. Only run-owned scratch is removed automatically. Completed checkpoints and requested outputs are retained.

Checkpoint and model-resolution parameters

• resume — Reuses valid completed milestones from the matching checkpoint namespace. Default: True. Missing, incomplete, corrupt, or incompatible milestones are recomputed. False creates a fresh namespace without deleting saved progress. Milestone writing remains enabled.
• checkpoint_dir — Shared root for persistent progress. Default: None, which uses <absolute FASTA path>.genatator-checkpoints. The runtime creates a fingerprint-specific directory inside it. This directory must be writable and accessible from every participating node.
• model_revisions — Optional mapping from edge, region, transcript_type, and segmentation to Hugging Face revisions. Default: None. Repository models are resolved to snapshots once by the coordinator. All workers load the resolved paths. Local model directories are used directly. Model, tokenizer, and code file hashes form part of checkpoint identity.
• local_files_only — Restricts model resolution to local paths and cached files. Default: False. Set it to True when the models are already available on shared storage and the allocation must not download checkpoints.

Reliability, reports, and cluster communication

• max_task_retries — Number of retries allowed for a failed work item, using the same inputs. Default: 1. A value of 0 disables retries. GPU out-of-memory microbatch reduction is handled separately before a task failure is reported.
• worker_timeout_seconds — Timeout for missing worker or node heartbeats. Default: 300.0. This is not a maximum allowed duration for a chromosome or model forward pass.
• startup_timeout_seconds — Maximum wait for worker/node startup and coordinator rendezvous. Default: 900.0.
• summary_path — Optional JSON inference-report path. Default: <output GFF path>.summary.json. Discovery-only runs use <FASTA path>.summary.json. The text report uses the same path with its suffix replaced by .txt. The Hugging Face interface also exposes the report as pipe.last_summary.
• reference_gff_path — Optional GFF for aggregate output validation. Default: None. A completed run passes when final transcript count and cumulative full genomic transcript length, including introns, match exactly. Probabilities, identifiers, file order, and individual boundary equality are not part of this aggregate validation. Output and reference paths must differ.
• coordinator_host — Reachable coordinator hostname/address advertised to Slurm node agents. Default: None, which uses the rank-zero node's fully qualified hostname. Set this when the compute network requires a different reachable address. Single-node execution uses a loopback listener.
• coordinator_port — Coordinator TCP port. Default: 0, which requests an available port. Specify a fixed allowed port when required by cluster networking rules.

What the pipeline does

1. Interval discovery

The first stage identifies candidate transcript intervals with two strand-aware DNA language models.

• The edge model detects transcription start site and polyadenylation signals, abbreviated as TSS and PolyA.
• The region model predicts intragenic signal, which is used to filter candidate intervals.

Edge prediction is processed in global chunks controlled by gene_finding_global_chunk_size. After each global chunk is peak-called, the raw edge predictions are discarded and only sparse peak coordinates are kept.

Region predictions are projected to nucleotide coordinates and reconstructed on CPU. The parallel reconstruction materializes the two intragenic channels needed for filtering, performs overlap and reverse-complement averaging, then thresholds them into Boolean strand-specific masks. Dense chromosome-length reconstruction arrays can still be required before thresholding, so RAM admission and optional temporary offloading apply to this stage.

Candidate intervals are formed by pairing strand-compatible TSS and PolyA peaks. Intervals with too much non-intragenic sequence are removed before transcript-type classification and segmentation.

The edge and region model calls may execute concurrently on different GPUs, but the CPU uses region results only after edge reconstruction, smoothing, peak calling, and candidate pairing. A chromosome with no candidates, or with no intervals left after region filtering, finishes with a valid empty result. Transcript-type inference, segmentation, and CDS processing are skipped for that chromosome only.

2. Transcript-type assignment

Each retained interval is classified by the transcript-type model as either mRNA or lnc_RNA. Only the leading token prefix defined by transcript_type_context_length is evaluated.

When reverse-complement averaging is enabled for this stage, forward and reverse-complement predictions are averaged. The reverse complement is formed from the whole interval before its leading token prefix is selected. The final decision is controlled by transcript_type_threshold.

3. Segmentation

Each retained interval is segmented into nucleotide-level structural classes by the segmentation model. Exons are derived from exon-versus-intron competition, and CDS segments are derived from CDS-versus-non-CDS competition.

Segmentation is stitched from interval blocks with zero regular overlap. The final block is shifted back when necessary to cover the interval end. Any overlapping positions are averaged. When tokenizer offset mappings are available, token-level outputs are projected to nucleotide coordinates. Reverse-complement segmentation, when enabled, processes the whole reversed interval and maps the resulting tracks back to the forward coordinate system before averaging.

Transcript-type and segmentation inference may execute concurrently after the chromosome's filtered interval list is complete. CPU annotation consumes intervals incrementally, in their logical order, once both results are available. CDS heuristics run only for mRNA transcripts when use_cds_heuristic=True.

4. GFF generation

The final annotation contains these feature types.

• gene
• mRNA or lnc_RNA
• exon
• CDS, for mRNA transcripts only

No CDS is emitted for lnc_RNA transcripts.

The GFF transcript attributes include lncRNA_probability, mRNA_probability, exon_segmentation_confidence, cds_segmentation_confidence, segmentation_confidence, and color. Exon and CDS features include mean_probability, and intron features are not emitted in the output GFF.

GPU scheduling and batching

Chromosome-stage assignments

Each GPU worker holds one stage model at a time. Its assignment identifies a chromosome, a model, and a set of windows or intervals. It processes that assignment without switching models between ordinary batches, transfers each completed batch's output to CPU, then moves the model off GPU when its assignment finishes. A temporary CPU-processing barrier does not end an unfinished chromosome-stage assignment.

CPU model and tokenizer caches remain available in each persistent GPU worker process. They are populated as models are needed and can hold multiple models on CPU. These are per-process copies, not one shared model object for the entire node, and their memory is included in node admission estimates. No stage-model forward pass runs on CPU.

CPU input preparation

Each spawned GPU process has one model-execution lane and a bounded pool of CPU preparation threads. There are no nested DataLoader processes and no fork calls inside a CUDA worker. The preparation threads read the assigned FASTA slices, apply the existing sequence orientation, tokenize with the model's supplied tokenizer, and assemble model-input rows. Each thread has an independent tokenizer instance so that concurrent padding and truncation settings cannot interfere.

The GPU execution lane retains tokenizer-based batch collation, model calls, output activation, and the blocking transfer of predictions to CPU. CPU threads can prepare upcoming inputs while that lane runs the current batch. They never execute a model, create CUDA tensors, or change a chromosome's GPU assignment. Completed preparations are matched by the original item identities before inference.

For edge and region, lookahead stays within the scheduler's released and memory-reserved window prefix. It does not bypass edge_prefetch_chunks or claim a second running GPU batch. If a helper GPU receives some pending work, an unused preparation hint may be discarded. For segmentation, whole intervals remain assigned to their GPU and only the existing internal window inputs are prepared ahead locally. RMT forward calls still have batch size 1.

Use these execution settings to control CPU support:

cpu_batch_workers_per_gpu=4,
cpu_prefetch_batches=2,
cpu_prefetch_buffer_gib_per_gpu=0.25,
cpu_threads_per_worker=1,

The equivalent command-line flags are --cpu-batch-workers-per-gpu, --cpu-prefetch-batches, --cpu-prefetch-buffer-gib-per-gpu, and --cpu-threads-per-worker. Setting cpu_batch_workers_per_gpu=0 disables threaded preparation without changing inference inputs.

Allocate CPUs to the node-agent task, not just to its launcher shell. For the four-GPU Slurm example, --cpus-per-task=32 provides room for four preparation threads per GPU and the other CPU work. The runtime never expands the inherited affinity mask. A job bound to one CPU cannot gain four CPU cores merely by creating four threads.

Native Rayon, OpenMP, MKL, and OpenBLAS thread pools are constrained in GPU subprocesses to avoid multiplying thread counts across the node. Fast Rust-backed tokenizers can run encoding work outside Python's interpreter lock. A Python-heavy tokenizer may gain mainly from overlapped I/O and GPU execution, so increasing the thread count is not guaranteed to improve every stage.

The summary records requested and effective preparation threads, preparation time summed over threads, input-wait time, prefetched input hits, discarded hints, and peak cache reservations. These timings distinguish waiting for model inputs from time spent in the model. Preparation time summed across threads is not an elapsed-time measurement.

Ready work for active chromosomes takes priority over starting another chromosome. Initially, independent chromosomes receive GPU work before helpers are assigned. When edge is active and region has not started, an available GPU starts region before joining edge as a helper. Edge and region have a small scheduling preference. Waiting type and segmentation jobs gain priority when needed to avoid indefinite delay.

When no eligible active-chromosome job can use an idle GPU, a pending chromosome may start. A not-yet-started stage can use otherwise spare GPUs from the outset. A helper may join an already-started stage only when no other pending or active chromosome has work left, including chromosomes temporarily waiting on CPU. Only unstarted items are redistributed. Running items are never interrupted or moved.

Work is divided by equal counts, with at most one item of difference between shares. Edge and region distribute their planned window inputs. Transcript type and segmentation use seeded random equal-count ownership of whole assembled intervals. A segmentation interval's internal windows stay on its assigned GPU. Random ownership does not reorder CPU reconstruction or annotation.

Each chromosome has its own completion and cancellation state. An empty or failed chromosome does not cancel unrelated chromosomes. GPU utilization can still be limited by CPU work, model loading, memory backpressure, or a final indivisible running interval.

Tokenization, padding, and effective batch sizes

Each model uses its supplied tokenizer. Edge and region rows use fixed token contexts with attention masks and offset mappings. Sequence classification and segmentation keep their own tokenization and special-token rules. The tokenizer performs batch padding. Batch-added padding is removed from token-level outputs before projection and reconstruction. Edge and region fixed-context trimming remains part of their projection procedure.

For Caduceus/Mamba models, models whose forward signature does not accept an attention mask, and left-padded tokenizers, the executor groups rows by exact encoded length. This avoids introducing padding into valid-sequence computations that do not use a padding mask. Consequently, the requested batch size of 32 is a maximum: heterogeneous intervals can form much smaller batches. A checkpoint's own forward implementation also determines whether it processes batch rows together or internally loops over them.

GPU out-of-memory recovery retries the same inputs in smaller microbatches, prints a warning, and records the reduction. The effective cap is reused for that model on the same worker. Context lengths, genomic windows, precision, and annotation thresholds are not reduced. An individual input that cannot fit at batch size 1 produces a reported failure.

Custom checkpoint interfaces must match their task: transcript type returns batch-first sequence scores, while the token-classification adapters expect batch-first token-aligned predictions. A segmentation model that directly returns a different nucleotide-output length requires a model-specific adapter. Padding removal alone does not provide that conversion.

Long chromosomes and memory management

The FASTA reader supports wrapped and unwrapped records and reads assigned sequence ranges without copying a complete chromosome into every GPU process. GPU memory is used for the active model, current inputs, activations, and outputs. Chromosome-scale tracks live in CPU memory or explicitly enabled temporary offload storage.

Completed GPU outputs are transferred to CPU after each forward call. Result credits limit dispatched work and protect the next results needed by the ordered CPU consumer. result_buffer_gib bounds coordinator result credits, while node admission separately accounts for reconstruction workspaces and receiver buffers. It is not necessary to retain all GPU outputs until a chromosome finishes.

For edge inference, each global chunk is reconstructed, averaged, smoothed, and peak-called before its dense arrays are released. edge_prefetch_chunks=1 permits GPU work for one following chunk to overlap CPU processing. A value of 0 disables that lookahead. Global chunk boundaries, overlapping model windows, and strand averaging are determined by the biological configuration, not GPU count.

Region reconstruction may require chromosome-length float arrays before Boolean thresholding. Single-window thresholding is not used as a substitute for averaging. Memory pressure limits concurrent chromosomes or enables temporary mapped arrays only when a positive disk budget is supplied. FFT operations, final masks, model caches, and third-party allocations still need RAM even when disk offloading is enabled.

A chromosome whose required workspace cannot fit any available node's budget is reported as a resource failure rather than processed with different chunking. RAM limits are conservative admission estimates, not hard process-memory guarantees. Leave headroom for allocators, transfer copies, and filesystem cache, and inspect recorded memory peaks.

Default model repositories

• edge_model_path, AIRI-Institute/genatator-moderngena-base-multispecies-edge-model
• region_model_path, AIRI-Institute/genatator-moderngena-base-multispecies-region-model
• transcript_type_model_path, AIRI-Institute/genatator-caduceus-ps-multispecies-transcript-type
• segmentation_model_path, AIRI-Institute/genatator-caduceus-ps-multispecies-segmentation

Input and output

Input

• Path to a FASTA file.
• The FASTA file may contain one record or multiple records.
• Multi-node runs require every participating node to access the same input, repository, resolved model paths, and checkpoint directory.

Output

• A single Python string, the path to the written .gff file.
• The file contents follow the GFF3 specification.
• Parallel runs also write JSON and text inference reports and persistent resume checkpoints.
• With predict_internal_structure=False, the call returns an empty string and writes no GFF. Use interval checkpoints and optional .bed files.
• A run with failed chromosomes can return a partial GFF. Check pipe.last_summary["complete"] or the command-line exit status before treating it as a complete annotation.

Checkpointing and resumption

resume=True is the default. The pipeline writes persistent checkpoints for completed units of biological work and reuses compatible completed milestones when launched with the same input and settings. Resumption occurs on an explicit launch. The pipeline does not automatically submit or requeue Slurm jobs.

A milestone is published only after its metadata and arrays are written and checksummed. Missing, incomplete, or corrupt results are not treated as completed work. Partial model state and partially evaluated intervals are not restored. Complete interval-type and segmentation results can be reused independently.

The checkpoint namespace is selected from FASTA content, biological settings, relevant source hashes, model/tokenizer/code files, and resolved labels. Changing GPU count, batch size, resource limits, or execution seed does not invalidate completed biological milestones. Each attempt records its execution settings so resumed output remains auditable. Changing biological identity selects a different namespace.

The shared filesystem must support cross-node POSIX file locking, atomic rename, and synchronization of writes. A run lock prevents simultaneous coordinators from publishing into one namespace, and a generation token fences writers belonging to a superseded attempt. Final output is rebuilt from completed chromosome annotations and finalized as a whole. The GFF is not appended during recovery.

resume=False or --no-resume creates a unique fresh namespace and still writes completed milestones. Automatic reuse with resume=True targets the normal matching fingerprint namespace, not a fresh namespace created with reuse disabled.

Persistent results and temporary storage

Checkpoint writing is independent of disk_offload_limit_gib_total. With temporary offloading disabled, completed milestone files are still written to checkpoint_dir. Allow enough shared disk capacity for segmentation-track checkpoints as well as metadata. The temporary offload budget does not limit persistent checkpoint size.

Completed checkpoints, requested intermediate files, final GFFs, and reports are retained. Normal finalization removes run-owned temporary offload files. A hard-killed process cannot perform cleanup. The next launch cleans incomplete writes and owned scratch for the matching run namespace. Cleanup does not remove unrelated runs' files. Cleanup outcomes and any failures are recorded in the report.

Dependencies

Create the Conda environment from environment.yml before running the pipeline locally. This project currently requires a CUDA-capable GPU.

conda env create -f environment.yml
conda activate genatator_pipeline

If the simple setup fails, use the robust staged setup. This follows the same strategy as Docker startup.

conda env create -n genatator_pipeline -f docker/conda-core.yml
conda activate genatator_pipeline

pip install torch==2.2.2+cu121 torchvision==0.17.2+cu121 torchaudio==2.2.2+cu121 --index-url https://download.pytorch.org/whl/cu121
pip install causal-conv1d==1.4.0 --no-build-isolation
pip install mamba-ssm==2.2.2 --no-build-isolation
pip install packaging==26.0 ninja==1.13.0 psutil==7.2.2
pip install flash-attn==2.6.3 --no-build-isolation
pip install -r docker/requirements.txt

Multi-node inference uses the same Python/CUDA environment on every node, shared filesystem access, and node-to-node TCP communication. The runtime uses Python multiprocessing and dependencies supplied by this repository. It does not require Ray, Dask, MPI, or a distributed-training backend. The CPU cache is not a CPU inference fallback.

Command-line inference

Single-node execution

Run from the repository directory in the configured environment:

python genatator_distributed.py \
  --fasta /shared/genome.fasta \
  --output-gff /shared/results/genome.gff \
  --devices 0 1 2 3 \
  --checkpoint-dir /shared/results/genatator-progress \
  --cpu-ram-limit-gib-per-node 100

Omitting --devices uses all visible GPUs on that node. The default biological configuration is config.json next to the launcher. Model paths and annotation settings can be supplied through another --config file. Execution-only settings can be supplied in JSON:

python genatator_distributed.py \
  --fasta /shared/genome.fasta \
  --output-gff /shared/results/genome.gff \
  --config config.json \
  --runtime-config examples/runtime.json \
  --checkpoint-dir /shared/results/genatator-progress

--runtime-config accepts runtime keys only. Unknown runtime keys are rejected. Explicit flags take precedence over the configuration files. Boolean flags include --resume / --no-resume, --local-files-only / --no-local-files-only, and --segmentation-is-rmt / --no-segmentation-is-rmt.

Temporary disk offloading requires explicit permission, for example:

python genatator_distributed.py \
  --fasta /shared/genome.fasta \
  --output-gff /shared/results/genome.gff \
  --devices 0 1 2 3 \
  --cpu-ram-limit-gib-per-node 100 \
  --disk-offload-limit-gib-total 200 \
  --offload-dir /shared/scratch/genatator

The 200 GiB budget applies across the run, not separately to each GPU or node. Rerun a command with the same checkpoint directory and default resume=True to recover completed progress. --no-resume starts a fresh namespace.

Use python genatator_distributed.py --help for the launch flags. --reference-mode selects the single-node, single-GPU diagnostic execution mode. It cannot be combined with --slurm.

Slurm multi-node inference

The launch layout is one Slurm task per node, with one spawned GPU worker per selected local GPU. Do not launch one node agent per GPU. The rank-zero task hosts the coordinator. Chromosome CPU owners may run on any participating node.

The supplied examples/slurm_multinode.sh requests two nodes, four GPUs per node, 16 CPUs per task, and --mem=120G per node, with a 100 GiB managed RAM budget. Adjust partition, account, time, CPU, GPU, and memory settings for the allocation. Activate the environment before submission or add the appropriate activation commands to the script.

export GENATATOR_REPO=/shared/genatator-pipeline
export GENATATOR_FASTA=/shared/genome.fasta
export GENATATOR_OUTPUT=/shared/results/genome.gff
export GENATATOR_CHECKPOINTS=/shared/results/genatator-progress

sbatch examples/slurm_multinode.sh

All nodes must see the repository, FASTA, model snapshots, and progress directory at the same paths. Repository snapshots are resolved once by the coordinator. Use a shared Hugging Face cache or explicit shared local model paths. A checkpoint available only on one node's local filesystem is not sufficient.

For an allocation that is already active, the corresponding node-agent launch is:

srun --ntasks-per-node=1 --cpus-per-task=32 --kill-on-bad-exit=0 --wait=0 \
  python -u genatator_distributed.py --slurm \
  --fasta /shared/genome.fasta \
  --output-gff /shared/results/genome.gff \
  --checkpoint-dir /shared/results/genatator-progress \
  --devices 0 1 2 3 \
  --cpu-ram-limit-gib-per-node 100

The allocation's total task count must equal its node count. The same visible-device list is checked on every node. The template also uses #SBATCH --no-kill. Worker and node failure behavior remains subject to site Slurm policy.

--coordinator-host selects a reachable address when the default hostname is unsuitable. --coordinator-port selects an allowed fixed TCP port. The default 0 requests an available port. --rendezvous optionally selects the shared launch-descriptor path. A unique GENATATOR_LAUNCH_ID, set by the template, identifies the launch.

The launch descriptor contains an RPC authentication key and is restricted to its owner. Communication uses authenticated Python Manager RPC, not encrypted transport. Only trusted job processes should reach the listener. Do not expose it to untrusted clients or outside the allocation.

A worker failure can be retried and isolated from other chromosomes. A lost chromosome CPU owner or coordinator is not transparently reconstructed in place. Completed milestones remain available for explicit relaunch in another allocation, and no automatic Slurm resubmission is performed.

Output annotation

The written GFF file contains one gene feature for each predicted gene locus and one transcript feature for each predicted transcript. Exons and CDS features are derived from the segmentation stage, and CDS features are emitted only for transcripts classified as mRNA.

The attribute field of each transcript includes transcript-type probabilities and segmentation-confidence values. The lncRNA_probability attribute stores the score produced by the transcript-type model.

Inference summary and output checks

Parallel inference writes <output GFF path>.summary.json and <output GFF path>.summary.txt by default. summary_path selects a different JSON location. Its text companion uses the .txt suffix. Discovery-only reports are placed next to the FASTA when no custom path is supplied. The Hugging Face call returns the GFF path, while pipe.last_summary contains the structured report.

Each checkpoint attempt also records run.json, events.jsonl, and a completed/partial summary.json. A handled fatal interruption records interrupted.json when the process is able to write it.

Stage/work-accounting details and milestone events are available through the runtime snapshot and event log. complete=False means the annotation is partial. An empty chromosome is a valid outcome. Failed inference is not converted into an empty success. Check the recorded reasons before accepting the output.

The distributed command returns exit code 0 when the run is complete and any requested comparison passes. Exit code 2 means a partial run or a failed/incomplete aggregate comparison. Fatal launch or runtime errors return a nonzero status. A Python caller should inspect pipe.last_summary["complete"] and pipe.last_summary["reference_comparison"]["status"].

Transcript-count and length validation

A reference comparison checks exactly two totals on the final GFF transcript features:

• Number of emitted mRNA and lnc_RNA features.
• Sum of their full genomic spans, including introns: end - start + 1 in GFF coordinates.

Each retained isoform contributes separately. Overlapping spans are not merged. These totals are measured after final filtering and deduplication. Matching totals is not a claim that every transcript boundary, probability, identifier, or GFF byte is identical.

Pass reference_gff_path to the Hugging Face call or --reference-gff-path to the launcher to include the check in the inference report:

python genatator_distributed.py \
  --fasta /shared/genome.fasta \
  --output-gff /shared/results/genome.gff \
  --devices 0 1 2 3 \
  --reference-gff-path /shared/validation/expected.gff

For two available annotation files:

python tools/compare_gff.py expected.gff results.gff --report comparison.json

The standalone comparator returns 0 for matching totals and 1 for a mismatch. In an inference report, not_checked means no reference was supplied. An incomplete result is never a passing comparison.

Docker deployment

All Docker assets are in docker/.

Build.

docker build -f docker/Dockerfile -t genatator-pipeline:latest .

Run.

docker run --gpus all --rm -p 3000:3000 -v "$(pwd)":/generated genatator-pipeline:latest

API endpoint.

• POST /api/genatator-pipeline/upload
• Input, multipart file containing FASTA, or form field dna.
• Output JSON fields, fasta_file, fai_file, gff_file, and archive.

Example.

curl -X POST "http://localhost:3000/api/genatator-pipeline/upload" -F "[email protected]"

The Flask application loads AIRI-Institute/genatator-pipeline and selects visible device 0. Its upload route accepts FASTA/DNA, not the complete set of distributed-runtime arguments. Use the command-line launcher for explicit GPU lists and Slurm allocations, or configure the pipeline construction in docker/app.py for the service's required defaults.

The API archive contains the input FASTA, its .fai index, and the GFF. Inference reports and resume checkpoints remain on the mounted output filesystem. They are not included in that archive. Keep /generated mounted persistently to retain those files across container lifetimes.