Toy GPT Training Repository¶

def argmax(values: list[float]) -> int:
    """Return the index of the maximum value in a list.

    Concept:
        argmax is the argument (index) at which a function reaches its maximum.

    Common uses:
        - Measuring accuracy during training (pick the most likely token)
        - Greedy decoding during inference (choose the top prediction)

    In training and inference:
        - A model outputs a probability distribution over possible next tokens.
        - The token with the highest probability is the model's most confident prediction.
        - argmax selects that token.

    Example:
        values = [0.1, 0.7, 0.2] has index values of 0,1, 2 respectively.
        argmax(values) -> 1 (since 0.7 is the largest value)

    Args:
        values: A non-empty list of numeric values (typically logits or probabilities).

    Returns:
        The index of the largest value in the list.

    Raises:
        ValueError: If values is empty.
    """
    if not values:
        raise ValueError("argmax() requires a non-empty list")

    best_idx: int = 0
    best_val: float = values[0]

    for i in range(1, len(values)):
        v = values[i]
        if v > best_val:
            best_val = v
            best_idx = i

    return best_idx

def cross_entropy_loss(probs: list[float], target_id: int) -> float:
    """Compute cross-entropy loss for a single training example.

    Concept: Cross-Entropy Loss
        Cross-entropy measures how well a predicted probability distribution
        matches the true outcome.

        In next-token prediction:
        - The true distribution is "one-hot" which means we encode it as either 1 or 0:
            - Probability = 1.0 for the correct next token
            - Probability = 0.0 for all others
        - The model predicts a probability distribution over all tokens.

        Cross-entropy answers the question:
            "How well does the predicted probability distribution align with the true outcome?"

    Formula:
        loss = -log(p_correct)

        - If the model assigns high probability to the correct token,
          the loss is small.
        - If the probability is near zero, the loss is large.

    Numerical safety:
        log(0) is undefined, so we clamp probabilities to a small minimum
        (1e-12). This does not change learning behavior in practice,
        but prevents runtime errors.

    In training:
        - This loss value drives gradient descent.
        - Lower loss means better predictions.

    Args:
        probs: A probability distribution over the vocabulary (sums to 1.0).
        target_id: The integer ID of the correct next token.

    Returns:
        A non-negative floating-point loss value.
        - 0.0 means a perfect prediction
        - Larger values indicate worse predictions
    Raises:
        ValueError: If target_id is out of range for probs.
    """
    if target_id < 0 or target_id >= len(probs):
        raise ValueError(
            f"target_id out of range: target_id={target_id} len(probs)={len(probs)}"
        )

    p: float = probs[target_id]

    # Guard against log(0), which would produce -infinity
    p = max(p, 1e-12)

    return -math.log(p)

def main() -> None:
    """Demonstrate a forward pass of the simple unigram model."""
    # Local imports keep modules decoupled.
    from toy_gpt_train.a_tokenizer import SimpleTokenizer
    from toy_gpt_train.b_vocab import Vocabulary

    log_header(LOG, "Simple Next-Token Model Demo (Unigram - No Context)")

    # Step 1: Tokenize input text.
    tokenizer: SimpleTokenizer = SimpleTokenizer()
    tokens: list[str] = tokenizer.get_tokens()

    if not tokens:
        LOG.info("No tokens available for demonstration.")
        return

    # Step 2: Build vocabulary.
    vocab: Vocabulary = Vocabulary(tokens)

    # Step 3: Initialize model.
    model: SimpleNextTokenModel = SimpleNextTokenModel(vocab_size=vocab.vocab_size())

    # Step 4: Forward pass (unigram ignores input).
    probs: list[float] = model.forward()

    # Step 5: Inspect results.
    LOG.info("Unigram ignores input - same predictions for any context:")
    LOG.info("Output probabilities for next token:")
    for idx, prob in enumerate(probs):
        tok: str | None = vocab.get_id_token(idx)
        LOG.info(f"  {tok!r} (ID {idx}) -> {prob:.4f}")

def main() -> None:
    """Run a simple training demo end-to-end."""
    from toy_gpt_train.a_tokenizer import CORPUS_DIR, SimpleTokenizer
    from toy_gpt_train.b_vocab import Vocabulary

    log_header(LOG, "Training Demo: Unigram (Frequency-Based) Model")

    base_dir: Final[Path] = Path(__file__).resolve().parents[2]
    outputs_dir: Final[Path] = base_dir / "outputs"
    train_log_path: Final[Path] = outputs_dir / "train_log.csv"

    # Step 0: Identify the corpus file (single file rule).
    corpus_path: Path = find_single_corpus_file(CORPUS_DIR)

    # Step 1: Load and tokenize the corpus.
    tokenizer: SimpleTokenizer = SimpleTokenizer(corpus_path=corpus_path)
    tokens: list[str] = tokenizer.get_tokens()

    if not tokens:
        LOG.error("No tokens found. Check corpus file.")
        return

    # Step 2: Build vocabulary (maps tokens <-> integer IDs).
    vocab: Vocabulary = Vocabulary(tokens)

    # Step 3: Convert token strings to integer IDs for training.
    token_ids: list[int] = []
    for tok in tokens:
        tok_id: int | None = vocab.get_token_id(tok)
        if tok_id is None:
            LOG.error(f"Token not found in vocabulary: {tok!r}")
            return
        token_ids.append(tok_id)

    # Step 4: Create training targets (just the tokens themselves for unigram).
    targets: list[int] = make_training_targets(token_ids)
    LOG.info(f"Created {len(targets)} training targets.")

    # Step 5: Initialize model (unigram has only 1 row of weights).
    model: SimpleNextTokenModel = SimpleNextTokenModel(vocab_size=vocab.vocab_size())

    # Step 6: Train the model.
    learning_rate: float = 0.1
    epochs: int = 50

    history: list[dict[str, float]] = train_model(
        model=model,
        targets=targets,
        learning_rate=learning_rate,
        epochs=epochs,
    )

    # Step 7: Save training metrics for analysis.
    write_training_log(train_log_path, history)

    # Step 7b: Write inspectable artifacts for downstream use.
    write_artifacts(
        base_dir=base_dir,
        corpus_path=corpus_path,
        vocab=vocab,
        model=model,
        model_kind="unigram",
        learning_rate=learning_rate,
        epochs=epochs,
        row_labeler=row_labeler_unigram(vocab, vocab.vocab_size()),
    )

    # Step 8: Qualitative check - what does the model predict?
    probs: list[float] = model.forward()
    best_id: int = argmax(probs)
    best_tok: str | None = vocab.get_id_token(best_id)
    LOG.info(
        f"After training, most likely token (based on frequency) "
        f"is {best_tok!r} (ID: {best_id})."
    )

def make_training_targets(token_ids: list[int]) -> list[int]:
    """Extract training targets for unigram model.

    For unigram, we don't need (input, target) pairs because
    the model ignores input. We just need the list of all tokens
    that appear in the corpus - each one is a target to predict.

    Args:
        token_ids: Sequence of integer token IDs from the corpus.

    Returns:
        List of target token IDs (all tokens in corpus).

    Example:
        Token sequence "the cat sat" with IDs [3, 1, 2] produces:
        [3, 1, 2]
        Meaning: the model should learn to predict these tokens
        based on their frequency (3 appears once, 1 appears once, etc.)
    """
    return token_ids

def row_labeler_unigram(vocab: VocabularyLike, vocab_size: int) -> RowLabeler:
    """Map a unigram row index to a label.

    Unigram has only one row, labeled to indicate it's context-free.
    """
    _ = vocab  # unused - unigram doesn't label by token
    _ = vocab_size

    def label(row_idx: int) -> str:
        # Only one row in unigram - label it descriptively
        return "(no context)"

    return label

def train_model(
    model: "SimpleNextTokenModel",
    targets: list[int],
    learning_rate: float,
    epochs: int,
) -> list[dict[str, float]]:
    """Train the unigram model using gradient descent on softmax cross-entropy.

    Unigram training learns corpus frequencies. The model has a single row
    of weights that gets updated for every token in the corpus.

    Training proceeds in epochs (full passes through all tokens).
    For each token, we:
    1. Compute the model's predicted probabilities (forward pass).
    2. Measure how wrong the prediction was (loss).
    3. Adjust weights to reduce the loss (gradient descent).

    Args:
        model: The model to train (weights will be modified in place).
        targets: List of target token IDs from the corpus.
        learning_rate: Step size for gradient descent.
        epochs: Number of complete passes through the training data.

    Returns:
        List of per-epoch metrics dictionaries containing epoch number,
        average loss, and accuracy.
    """
    history: list[dict[str, float]] = []

    for epoch in range(1, epochs + 1):
        total_loss: float = 0.0
        correct: int = 0

        for target_id in targets:
            # Forward pass: get probability distribution (same for all inputs).
            probs: list[float] = model.forward()

            # Compute loss: how surprised is the model by this token?
            loss: float = cross_entropy_loss(probs, target_id)
            total_loss += loss

            # Check if the model's top prediction matches the target.
            pred_id: int = argmax(probs)
            if pred_id == target_id:
                correct += 1

            # Backward pass: update the single row of weights.
            #
            # For softmax cross-entropy, the gradient is:
            #   gradient[j] = predicted_prob[j] - true_prob[j]
            #
            # This pushes probability mass toward frequently-seen tokens.
            row: list[float] = model.weights[0]  # unigram has only one row
            for j in range(model.vocab_size):
                y: float = 1.0 if j == target_id else 0.0
                grad: float = probs[j] - y
                row[j] -= learning_rate * grad

        # Compute epoch-level metrics.
        avg_loss: float = total_loss / len(targets) if targets else float("nan")
        accuracy: float = correct / len(targets) if targets else 0.0

        metrics: dict[str, float] = {
            "epoch": float(epoch),
            "avg_loss": avg_loss,
            "accuracy": accuracy,
        }
        history.append(metrics)

        LOG.info(
            f"Epoch {epoch}/{epochs} | avg_loss={avg_loss:.6f} | accuracy={accuracy:.3f}"
        )

    return history

def generate_tokens_unigram(
    model: SimpleNextTokenModel,
    vocab: ArtifactVocabulary,
    num_tokens: int,
) -> list[str]:
    """Generate tokens using unigram (no context - same prediction every time).

    Note: Unigram ignores all context, so there's no start token.
    Every generated token will be the same (the most frequent word).
    """
    generated: list[str] = []

    for _ in range(num_tokens):
        probs: list[float] = model.forward()  # no input - unigram ignores context
        next_id: int = argmax(probs)
        next_token: str | None = vocab.get_id_token(next_id)

        if next_token is None:
            LOG.error(f"Generated invalid token ID: {next_id}")
            break

        generated.append(next_token)

    return generated

def load_meta(path: Path) -> JsonObject:
    """Load 00_meta.json."""
    with path.open("r", encoding="utf-8") as f:
        data: JsonObject = json.load(f)
    return data

def load_model_weights_csv(path: Path, vocab_size: int) -> list[list[float]]:
    """Load 02_model_weights.csv -> weights matrix.

    For unigram, expects exactly 1 row (context-independent predictions).
    """
    weights: list[list[float]] = []

    with path.open("r", encoding="utf-8", newline="") as f:
        reader = csv.reader(f)
        header = next(reader, None)
        if header is None:
            raise ValueError("Weights CSV is empty.")
        if len(header) < 2 or header[0] != "input_token":
            raise ValueError("Weights CSV must start with header 'input_token'.")

        num_outputs = len(header) - 1
        if num_outputs != vocab_size:
            raise ValueError(
                f"Weights CSV output width mismatch. Expected {vocab_size} output columns "
                f"but found {num_outputs}."
            )

        for row in reader:
            if not row:
                continue
            if len(row) != vocab_size + 1:
                raise ValueError(
                    f"Invalid weights row length. Expected {vocab_size + 1} columns but found {len(row)}."
                )
            # row[0] is input token label; row[1:] are numeric weights
            weights.append([float(x) for x in row[1:]])

    # Unigram has exactly 1 row
    if len(weights) != 1:
        raise ValueError(
            f"Unigram weights CSV should have 1 row but found {len(weights)}."
        )

    return weights

def load_vocabulary_csv(path: Path) -> ArtifactVocabulary:
    """Load 01_vocabulary.csv -> ArtifactVocabulary."""
    token_to_id: dict[str, int] = {}
    id_to_token: dict[int, str] = {}
    token_freq: dict[str, int] = {}

    with path.open("r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        expected = {"token_id", "token", "frequency"}
        if reader.fieldnames is None or set(reader.fieldnames) != expected:
            raise ValueError(
                f"Unexpected vocabulary header. Expected {sorted(expected)} "
                f"but got {reader.fieldnames}"
            )

        for row in reader:
            token_id = int(row["token_id"])
            token = row["token"]
            freq = int(row["frequency"])

            token_to_id[token] = token_id
            id_to_token[token_id] = token
            token_freq[token] = freq

    return ArtifactVocabulary(
        token_to_id=token_to_id,
        id_to_token=id_to_token,
        token_freq=token_freq,
    )

def main() -> None:
    """Run inference using saved training artifacts."""
    log_header(LOG, "Inference Demo: Load Artifacts and Generate Text (Unigram)")

    base_dir: Final[Path] = Path(__file__).resolve().parents[2]
    artifacts_dir: Final[Path] = base_dir / "artifacts"
    meta_path: Final[Path] = artifacts_dir / "00_meta.json"
    vocab_path: Final[Path] = artifacts_dir / "01_vocabulary.csv"
    weights_path: Final[Path] = artifacts_dir / "02_model_weights.csv"
    require_artifacts(
        meta_path=meta_path,
        vocab_path=vocab_path,
        weights_path=weights_path,
        train_hint="uv run python src/toy_gpt_train/d_train.py",
    )

    meta: JsonObject = load_meta(meta_path)
    vocab: ArtifactVocabulary = load_vocabulary_csv(vocab_path)

    v: int = vocab.vocab_size()
    model: SimpleNextTokenModel = SimpleNextTokenModel(vocab_size=v)
    model.weights = load_model_weights_csv(weights_path, vocab_size=v)

    args: argparse.Namespace = parse_args([])

    LOG.info(
        f"Loaded repo_name={meta.get('repo_name')} model_kind={meta.get('model_kind')}"
    )
    LOG.info(f"Vocab size: {v}")
    LOG.info("Unigram model: predictions are the same regardless of input.")

    # Show predictions (same for any input)
    probs: list[float] = model.forward()
    LOG.info("Top next-token predictions (based on corpus frequency):")
    for tok_id, prob in top_k(probs, k=max(1, args.topk)):
        tok: str | None = vocab.get_id_token(tok_id)
        LOG.info(f"  {tok!r} (ID {tok_id}): {prob:.4f}")

    generated: list[str] = generate_tokens_unigram(
        model=model,
        vocab=vocab,
        num_tokens=max(0, args.num_tokens),
    )

    LOG.info("Generated sequence:")
    LOG.info(f"  {' '.join(generated)}")

def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    """Parse command-line arguments for inference."""
    parser = argparse.ArgumentParser(
        description="Run inference using saved training artifacts (unigram)."
    )
    parser.add_argument(
        "--num-tokens",
        type=int,
        default=10,
        help="Number of tokens to generate (default: 10).",
    )
    parser.add_argument(
        "--topk",
        type=int,
        default=3,
        help="Number of top predictions to display (default: 3).",
    )
    return parser.parse_args(argv)

def require_artifacts(
    *,
    meta_path: Path,
    vocab_path: Path,
    weights_path: Path,
    train_hint: str,
) -> None:
    """Fail fast with a helpful message if artifacts are missing."""
    missing: list[Path] = []
    for p in [meta_path, vocab_path, weights_path]:
        if not p.exists():
            missing.append(p)

    if missing:
        LOG.error("Missing training artifacts:")
        for p in missing:
            LOG.error(f"  - {p}")
        LOG.error("Run training first:")
        LOG.error(f"  {train_hint}")
        raise SystemExit(2)

def top_k(probs: list[float], k: int) -> list[tuple[int, float]]:
    """Return top-k (token_id, probability) pairs sorted by probability."""
    pairs: list[tuple[int, float]] = list(enumerate(probs))
    pairs.sort(key=lambda x: x[1], reverse=True)
    return pairs[:k]

def artifact_paths_from_base_dir(base_dir: Path) -> dict[str, Path]:
    """Return standard artifact paths under base_dir/artifacts/."""
    artifacts_dir = artifacts_dir_from_base_dir(base_dir)
    return {
        "00_meta.json": artifacts_dir / "00_meta.json",
        "01_vocabulary.csv": artifacts_dir / "01_vocabulary.csv",
        "02_model_weights.csv": artifacts_dir / "02_model_weights.csv",
        "03_token_embeddings.csv": artifacts_dir / "03_token_embeddings.csv",
    }

def artifacts_dir_from_base_dir(base_dir: Path) -> Path:
    """Return artifacts/ directory under a repository base directory."""
    return base_dir / "artifacts"

def find_single_corpus_file(corpus_dir: Path) -> Path:
    """Find the single corpus file in corpus/ (same rule as SimpleTokenizer)."""
    if not corpus_dir.exists():
        msg = f"Corpus directory not found: {corpus_dir}"
        raise FileNotFoundError(msg)

    files = sorted([p for p in corpus_dir.iterdir() if p.is_file()])
    if len(files) == 0:
        msg = f"No files found in corpus directory: {corpus_dir}"
        raise FileNotFoundError(msg)
    if len(files) > 1:
        msg = f"Expected exactly one file in corpus directory, found {len(files)}: {corpus_dir}"
        raise ValueError(msg)

    return files[0]

def outputs_dir_from_base_dir(base_dir: Path) -> Path:
    """Return outputs/ directory under a repository base directory."""
    return base_dir / "outputs"

def repo_name_from_base_dir(base_dir: Path) -> str:
    """Infer repository name from base directory."""
    return base_dir.resolve().name

def sha256_of_bytes(data: bytes) -> str:
    """Return hex SHA-256 digest for given bytes."""
    return hashlib.sha256(data).hexdigest()

def sha256_of_file(path: Path) -> str:
    """Return hex SHA-256 digest for a file."""
    data = path.read_bytes()
    return sha256_of_bytes(data)

def write_artifacts(
    *,
    base_dir: Path,
    corpus_path: Path,
    vocab: VocabularyLike,
    model: SimpleNextTokenModel,
    model_kind: str,
    learning_rate: float,
    epochs: int,
    row_labeler: RowLabeler,
) -> None:
    """Write all training artifacts to artifacts/.

    Args:
        base_dir: Repository base directory.
        corpus_path: Corpus file used for training.
        vocab: VocabularyLike instance.
        model: Trained model (weights already updated).
        model_kind: Human-readable model kind (e.g., "unigram", "bigram").
        learning_rate: Training learning rate.
        epochs: Number of training passes.
        row_labeler: Function that maps a model weight-row index to a label
            written in the first column of 02_model_weights.csv.
    """
    artifacts_dir: Final[Path] = base_dir / "artifacts"
    meta_path: Final[Path] = artifacts_dir / "00_meta.json"
    vocab_path: Final[Path] = artifacts_dir / "01_vocabulary.csv"
    weights_path: Final[Path] = artifacts_dir / "02_model_weights.csv"
    embeddings_path: Final[Path] = artifacts_dir / "03_token_embeddings.csv"

    artifacts_dir.mkdir(parents=True, exist_ok=True)

    write_vocabulary_csv(vocab_path, vocab)
    write_model_weights_csv(weights_path, vocab, model, row_labeler=row_labeler)
    write_token_embeddings_csv(embeddings_path, model, row_labeler=row_labeler)
    write_meta_json(
        meta_path,
        base_dir=base_dir,
        corpus_path=corpus_path,
        vocab_size=vocab.vocab_size(),
        model_kind=model_kind,
        learning_rate=learning_rate,
        epochs=epochs,
    )