entrenar 0.7.10

Training & Optimization library with autograd, LoRA, quantization, and model merging
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

//! Batch struct for training

use ndarray::Array2;

/// Batch of examples for training
#[derive(Debug, Clone)]
pub struct Batch {
    /// Input IDs [batch_size, max_seq_len]
    pub input_ids: Array2<u32>,
    /// Attention mask [batch_size, max_seq_len]
    pub attention_mask: Array2<u8>,
    /// Labels [batch_size, max_seq_len] (optional)
    pub labels: Option<Array2<u32>>,
    /// Original sequence lengths
    pub lengths: Vec<usize>,
}

impl Batch {
    /// Get batch size
    #[must_use]
    pub fn batch_size(&self) -> usize {
        self.input_ids.nrows()
    }

    /// Get maximum sequence length
    #[must_use]
    pub fn max_seq_len(&self) -> usize {
        self.input_ids.ncols()
    }
}