Struct LMBatch 

pub struct LMBatch {
    pub input_ids: Vec<u32>,
    pub target_ids: Vec<u32>,
    pub batch_size: usize,
    pub seq_len: usize,
}

A batch of tokenized sequences for language model training

Fields

input_ids: Vec<u32>

Input token IDs (batch_size x seq_len flattened)

target_ids: Vec<u32>

Target token IDs (batch_size x seq_len flattened, shifted by 1)

batch_size: usize

Batch size

seq_len: usize

Sequence length

Implementations

impl LMBatch

pub fn from_sequences(sequences: &[Vec<u32>], pad_id: u32, eos_id: u32) -> Self

Create a new LM batch from token sequences

For causal LM, targets are inputs shifted by 1: input: [BOS, A, B, C, D] target: [A, B, C, D, EOS]

pub fn single(input_ids: Vec<u32>, target_ids: Vec<u32>) -> Self

Create a batch from a single sequence (for testing)

pub fn get_input(&self, batch_idx: usize) -> Option<&[u32]>

Get input IDs for a specific batch item

pub fn get_target(&self, batch_idx: usize) -> Option<&[u32]>

Get target IDs for a specific batch item

pub fn num_tokens(&self) -> usize

Total number of tokens in batch

Trait Implementations

impl Clone for LMBatch

fn clone(&self) -> LMBatch

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for LMBatch

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.