Как составлять авторегрессионные пары «вход — цель»

/// One task-specific row whose label had to be supplied separately.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct HandLabeledRow<'a> {
    /// Fixed input chosen by the example author.
    pub input: &'a [u32],
    /// Separate sentiment class attached to that input.
    pub label: SentimentLabel,
}

/// A human-supplied class for the synthetic sentiment contrast.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SentimentLabel {
    /// The example expresses negative sentiment.
    Negative,
    /// The example expresses positive sentiment.
    Positive,
}

impl SentimentLabel {
    /// Returns the label as display text.
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Negative => "negative",
            Self::Positive => "positive",
        }
    }
}

/// Returns two tiny task-specific rows with separately supplied labels.
pub const fn hand_labeled_rows() -> [HandLabeledRow<'static>; 2] {
    [
        HandLabeledRow {
            input: &[41, 42, 43],
            label: SentimentLabel::Negative,
        },
        HandLabeledRow {
            input: &[51, 52, 53],
            label: SentimentLabel::Positive,
        },
    ]
}

/// Positive sizes that determine which candidate starts may produce pairs.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct CausalWindowConfig {
    context_length: usize,
    stride: usize,
    required_source_tokens: usize,
}

impl CausalWindowConfig {
    /// Validates the context length and distance between candidate starts.
    pub const fn new(
        context_length: usize,
        stride: usize,
    ) -> Result<Self, CausalWindowConfigError> {
        if context_length == 0 {
            return Err(CausalWindowConfigError::ZeroContextLength);
        }
        if stride == 0 {
            return Err(CausalWindowConfigError::ZeroStride);
        }
        let Some(required_source_tokens) = context_length.checked_add(1) else {
            return Err(CausalWindowConfigError::ContextLengthOverflow);
        };
        Ok(Self {
            context_length,
            stride,
            required_source_tokens,
        })
    }

    /// Returns the number of input IDs and target IDs in every emitted pair.
    pub const fn context_length(self) -> usize {
        self.context_length
    }

    /// Returns the distance between consecutive candidate starts.
    pub const fn stride(self) -> usize {
        self.stride
    }

    /// Returns the number of source IDs required to emit one shifted pair.
    pub const fn required_source_tokens(self) -> usize {
        self.required_source_tokens
    }

    /// Counts complete pairs without adding potentially huge indices.
    pub const fn window_count(self, document_length: usize) -> usize {
        if document_length < self.required_source_tokens {
            0
        } else {
            (document_length - self.required_source_tokens) / self.stride + 1
        }
    }

    /// Borrows every complete shifted pair selected inside one token slice.
    pub const fn windows<'a>(self, tokens: &'a [u32]) -> CausalWindows<'a> {
        CausalWindows {
            tokens,
            config: self,
            next_start: Some(0),
            remaining: self.window_count(tokens.len()),
        }
    }

    /// Returns the suffix at the first candidate start that cannot fill a pair.
    ///
    /// `None` means that the next candidate start lies at or beyond the end of
    /// the document. A returned suffix may overlap earlier complete pairs.
    pub fn incomplete_tail<'a>(self, tokens: &'a [u32]) -> Option<IncompleteTail<'a>> {
        let window_count = self.window_count(tokens.len());
        let start = if window_count == 0 {
            0
        } else {
            (window_count - 1)
                .checked_mul(self.stride)?
                .checked_add(self.stride)?
        };
        (start < tokens.len()).then(|| IncompleteTail {
            start,
            tokens: &tokens[start..],
        })
    }
}
/// A rejected configuration for causal-example construction.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CausalWindowConfigError {
    /// An empty input and target would not teach a next-token relation.
    ZeroContextLength,
    /// `T + 1` cannot be represented, so the source requirement is undefined.
    ContextLengthOverflow,
    /// A zero stride would select the same start forever.
    ZeroStride,
}

impl fmt::Display for CausalWindowConfigError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::ZeroContextLength => "context length must be positive",
            Self::ContextLengthOverflow => {
                "context length is too large to require one additional source token"
            }
            Self::ZeroStride => "stride must be positive",
        })
    }
}

impl Error for CausalWindowConfigError {}

/// One complete input/target pair borrowed from a single document.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct CausalWindow<'a> {
    start: usize,
    input: &'a [u32],
    target: &'a [u32],
}

impl<'a> CausalWindow<'a> {
    /// Returns the input's zero-based position inside its document.
    pub const fn start(self) -> usize {
        self.start
    }

    /// Returns exactly `context_length` source token IDs.
    pub const fn input(self) -> &'a [u32] {
        self.input
    }

    /// Returns the `T`-token source slice beginning one position after the input.
    pub const fn target(self) -> &'a [u32] {
        self.target
    }
}

/// The first selected suffix that is too short to emit a complete pair.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct IncompleteTail<'a> {
    start: usize,
    tokens: &'a [u32],
}

impl<'a> IncompleteTail<'a> {
    /// Returns the candidate position at which the too-short suffix begins.
    pub const fn start(self) -> usize {
        self.start
    }

    /// Returns the remaining source IDs; these may overlap earlier pairs.
    pub const fn tokens(self) -> &'a [u32] {
        self.tokens
    }
}

/// A repeatable, exact-size iterator over one document's complete pairs.
#[derive(Clone, Debug)]
pub struct CausalWindows<'a> {
    tokens: &'a [u32],
    config: CausalWindowConfig,
    next_start: Option<usize>,
    remaining: usize,
}

impl<'a> Iterator for CausalWindows<'a> {
    type Item = CausalWindow<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.remaining == 0 {
            return None;
        }

        let start = self.next_start?;
        let input_end = start.checked_add(self.config.context_length)?;
        let target_start = start.checked_add(1)?;
        let target_end = input_end.checked_add(1)?;
        let input = self.tokens.get(start..input_end)?;
        let target = self.tokens.get(target_start..target_end)?;

        self.remaining -= 1;
        self.next_start = if self.remaining == 0 {
            None
        } else {
            start.checked_add(self.config.stride)
        };
        Some(CausalWindow {
            start,
            input,
            target,
        })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.remaining, Some(self.remaining))
    }
}

impl ExactSizeIterator for CausalWindows<'_> {}
impl FusedIterator for CausalWindows<'_> {}

/// One owned token sequence with its stable document and partition identities.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EncodedDocument {
    id: String,
    partition: Partition,
    token_ids: Vec<u32>,
}

impl EncodedDocument {
    /// Returns the stable source-document identity.
    pub fn id(&self) -> &str {
        &self.id
    }

    /// Returns the frozen partition role inherited from the split manifest.
    pub const fn partition(&self) -> Partition {
        self.partition
    }

    /// Returns the separately wrapped `[BOS, content..., EOS]` sequence.
    pub fn token_ids(&self) -> &[u32] {
        &self.token_ids
    }

    /// Opens a fresh borrowed pair iterator without consuming this document.
    pub fn windows(&self, config: CausalWindowConfig) -> CausalWindows<'_> {
        config.windows(&self.token_ids)
    }

    /// Reports the first candidate suffix that cannot fill a complete pair.
    pub fn incomplete_tail(&self, config: CausalWindowConfig) -> Option<IncompleteTail<'_>> {
        config.incomplete_tail(&self.token_ids)
    }
}

/// Encoded documents kept in three disjoint owned collections.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EncodedCorpusPartitions {
    train: Vec<EncodedDocument>,
    validation: Vec<EncodedDocument>,
    test: Vec<EncodedDocument>,
}

impl EncodedCorpusPartitions {
    /// Applies one frozen tokenizer independently to every frozen source document.
    pub fn from_partitions(partitions: &CorpusPartitions<'_>, tokenizer: &BpeTokenizer) -> Self {
        let encode = |partition| {
            partitions
                .documents(partition)
                .iter()
                .map(|document| EncodedDocument {
                    id: document.id().to_owned(),
                    partition,
                    token_ids: tokenizer.encode_utf8_document(document.text()),
                })
                .collect()
        };

        Self {
            train: encode(Partition::Train),
            validation: encode(Partition::Validation),
            test: encode(Partition::Test),
        }
    }

    /// Returns only the separately encoded documents for one partition.
    pub fn documents(&self, partition: Partition) -> &[EncodedDocument] {
        match partition {
            Partition::Train => &self.train,
            Partition::Validation => &self.validation,
            Partition::Test => &self.test,
        }
    }
}

use ch05_autoregressive_examples::hand_labeled_rows;
use llm_from_scratch::corpus::{Corpus, Partition, SplitManifest};
use llm_from_scratch::data::{CausalWindowConfig, EncodedCorpusPartitions};
use llm_from_scratch::tokenizer::bpe::BpeTokenizer;
use llm_from_scratch::tokenizer::bpe_trainer::BpeTrainer;

const CORPUS_JSON: &str = include_str!("../../../data/tiny-bilingual-corpus.json");
const SPLIT_MANIFEST: &str = include_str!("../../../data/splits.json");

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let rows = hand_labeled_rows();
    println!(
        "task-specific hand-labeled contrast: sentiment rows={}",
        rows.len()
    );
    for (index, row) in rows.iter().enumerate() {
        println!(
            "labeled row {index}: input={:?} label={}",
            row.input,
            row.label.as_str()
        );
    }

    let source = [0, 41, 42, 43, 44, 1];
    let config = CausalWindowConfig::new(3, 1)?;
    println!("source sequence for next-token targets: {source:?}");
    println!(
        "config: context={} stride={} required={}",
        config.context_length(),
        config.stride(),
        config.required_source_tokens()
    );
    println!("generated pairs: {}", config.window_count(source.len()));
    for window in config.windows(&source) {
        println!(
            "pair start={} input={:?} target={:?}",
            window.start(),
            window.input(),
            window.target()
        );
    }
    let tail = config
        .incomplete_tail(&source)
        .expect("worked source has one too-short suffix");
    println!(
        "next start too short: start={} tokens={:?} required={} emitted=false",
        tail.start(),
        tail.tokens(),
        config.required_source_tokens()
    );

    let short = [0, 61, 1];
    let short_tail = config
        .incomplete_tail(&short)
        .expect("short document leaves a too-short suffix");
    println!(
        "short document: {short:?} pairs={} suffix={:?}",
        config.window_count(short.len()),
        short_tail.tokens()
    );

    let corpus = Corpus::from_json(CORPUS_JSON)?;
    let manifest = SplitManifest::from_json(SPLIT_MANIFEST)?;
    let partitions = manifest.partition(&corpus)?;
    let training = BpeTrainer::new(8).train(&partitions)?;
    let tokenizer = BpeTokenizer::from_training(&training)?;
    let encoded = EncodedCorpusPartitions::from_partitions(&partitions, &tokenizer);
    println!(
        "frozen encoded documents: train={} validation={} test={}",
        encoded.documents(Partition::Train).len(),
        encoded.documents(Partition::Validation).len(),
        encoded.documents(Partition::Test).len()
    );
    let pair_counts = [Partition::Train, Partition::Validation, Partition::Test].map(|partition| {
        encoded
            .documents(partition)
            .iter()
            .map(|document| document.windows(config).count())
            .sum::<usize>()
    });
    println!(
        "encoded pairs at context=3 stride=1: train={} validation={} test={}",
        pair_counts[0], pair_counts[1], pair_counts[2]
    );

    println!("chapter 6 handoff: count each adjacent training-document transition once");
    Ok(())
}