setup

from setuptools import setup, find_packages

setup(
    name="bigram-lm",
    version="0.1.0",
    packages=find_packages(),
    install_requires=[
        "torch>=2.0.0",
        "numpy>=1.21.0",
    ],
    extras_require={
        'test': [
            'pytest>=7.0.0',
            'pytest-cov>=4.0.0',
        ],
    },
    author="Your Name",
    author_email="your.email@example.com",
    description="A simple bigram language model implementation",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",
    python_requires=">=3.7",
)


tests.conftest

import pytest
import torch
from src.model import BigramLanguageModel
from src.data import TextDataset

@pytest.fixture
def sample_text():
    return "Hello, World!"

@pytest.fixture
def dataset(sample_text):
    return TextDataset(sample_text, sequence_length=3)

@pytest.fixture
def model(dataset):
    return BigramLanguageModel(dataset.vocab_size)

@pytest.fixture
def device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'


import torch
import pytest
from src.model import BigramLanguageModel

def test_model_initialization(model, dataset):
    assert isinstance(model, BigramLanguageModel)
    assert model.bigram_table.shape == (dataset.vocab_size, dataset.vocab_size)

def test_model_forward(model, device):
    batch_size, seq_length = 2, 4
    idx = torch.randint(0, model.vocab_size, (batch_size, seq_length)).to(device)
    output = model(idx)

    assert output.shape == (batch_size, seq_length, model.vocab_size)
    assert not torch.isnan(output).any()

def test_model_generate(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    max_new_tokens = 10

    generated = model.generate(context, max_new_tokens)
    assert generated.shape == (1, max_new_tokens + 1)
    assert not torch.isnan(generated).any()

def test_temperature_effect(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)

    # Higher temperature should lead to more randomness
    high_temp = model.generate(context, max_new_tokens=100, temperature=2.0)
    low_temp = model.generate(context, max_new_tokens=100, temperature=0.1)

    # Convert to sets to compare unique tokens
    high_temp_unique = len(set(high_temp[0].tolist()))
    low_temp_unique = len(set(low_temp[0].tolist()))

    assert high_temp_unique >= low_temp_unique

def test_bigram_properties(model):
    # Test that each row in bigram_table represents a probability distribution after softmax
    probs = torch.softmax(model.bigram_table, dim=1)
    row_sums = probs.sum(dim=1)
    assert torch.allclose(row_sums, torch.ones_like(row_sums))


import torch
import pytest
from src.model import BigramLanguageModel

def test_model_initialization(model, dataset):
    assert isinstance(model, BigramLanguageModel)
    assert model.bigram_table.shape == (dataset.vocab_size, dataset.vocab_size)

def test_model_forward(model, device):
    batch_size, seq_length = 2, 4
    idx = torch.randint(0, model.vocab_size, (batch_size, seq_length)).to(device)
    output = model(idx)

    assert output.shape == (batch_size, seq_length, model.vocab_size)
    assert not torch.isnan(output).any()

def test_model_generate(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    max_new_tokens = 10

    generated = model.generate(context, max_new_tokens)
    assert generated.shape == (1, max_new_tokens + 1)
    assert not torch.isnan(generated).any()

def test_temperature_effect(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)


import torch
import pytest
from src.model import BigramLanguageModel

def test_model_initialization(model, dataset):
    assert isinstance(model, BigramLanguageModel)
    assert model.bigram_table.shape == (dataset.vocab_size, dataset.vocab_size)

def test_model_forward(model, device):
    batch_size, seq_length = 2, 4
    idx = torch.randint(0, model.vocab_size, (batch_size, seq_length)).to(device)
    output = model(idx)

    assert output.shape == (batch_size, seq_length, model.vocab_size)
    assert not torch.isnan(output).any()

def test_model_generate(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    max_new_tokens = 10

    generated = model.generate(context, max_new_tokens)
    assert generated.shape == (1, max_new_tokens + 1)
    assert not torch.isnan(generated).any()

def test_temperature_effect(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)

    # Higher temperature should lead to more randomness
    high_temp = model.generate(context, max_new_tokens=100, temperature=2.0)
    low_temp = model.generate(context, max_new_tokens=100, temperature=0.1)

    # Convert to sets to compare unique tokens
    high_temp_unique = len(set(high_temp[0].tolist()))
    low_temp_unique = len(set(low_temp[0].tolist()))


def test_model_initialization(model, dataset):
    assert isinstance(model, BigramLanguageModel)
    assert model.bigram_table.shape == (dataset.vocab_size, dataset.vocab_size)

def test_model_forward(model, device):
    batch_size, seq_length = 2, 4
    idx = torch.randint(0, model.vocab_size, (batch_size, seq_length)).to(device)
    output = model(idx)

    assert output.shape == (batch_size, seq_length, model.vocab_size)
    assert not torch.isnan(output).any()

def test_model_generate(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    max_new_tokens = 10

    generated = model.generate(context, max_new_tokens)
    assert generated.shape == (1, max_new_tokens + 1)
    assert not torch.isnan(generated).any()

def test_temperature_effect(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)

    # Higher temperature should lead to more randomness
    high_temp = model.generate(context, max_new_tokens=100, temperature=2.0)
    low_temp = model.generate(context, max_new_tokens=100, temperature=0.1)

    # Convert to sets to compare unique tokens
    high_temp_unique = len(set(high_temp[0].tolist()))
    low_temp_unique = len(set(low_temp[0].tolist()))

    assert high_temp_unique >= low_temp_unique

def test_bigram_properties(model):
    # Test that each row in bigram_table represents a probability distribution after softmax
    probs = torch.softmax(model.bigram_table, dim=1)
    row_sums = probs.sum(dim=1)
    assert torch.allclose(row_sums, torch.ones_like(row_sums))


    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    max_new_tokens = 10

    generated = model.generate(context, max_new_tokens)
    assert generated.shape == (1, max_new_tokens + 1)
    assert not torch.isnan(generated).any()

def test_temperature_effect(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)

    # Higher temperature should lead to more randomness
    high_temp = model.generate(context, max_new_tokens=100, temperature=2.0)
    low_temp = model.generate(context, max_new_tokens=100, temperature=0.1)

    # Convert to sets to compare unique tokens
    high_temp_unique = len(set(high_temp[0].tolist()))
    low_temp_unique = len(set(low_temp[0].tolist()))

    assert high_temp_unique >= low_temp_unique

def test_bigram_properties(model):
    # Test that each row in bigram_table represents a probability distribution after softmax
    probs = torch.softmax(model.bigram_table, dim=1)
    row_sums = probs.sum(dim=1)
    assert torch.allclose(row_sums, torch.ones_like(row_sums))


import torch
from src.model import BigramLanguageModel
from src.data import TextDataset

def test_end_to_end_generation(sample_text, device):
    # Create dataset and model
    dataset = TextDataset(sample_text, sequence_length=3)
    model = BigramLanguageModel(dataset.vocab_size).to(device)

    # Generate text
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    generated_indices = model.generate(context, max_new_tokens=20)
    generated_text = dataset.decode(generated_indices[0].tolist())

    assert len(generated_text) == 21  # context + max_new_tokens
    assert all(c in dataset.stoi for c in generated_text)

def test_generation_with_different_temperatures(sample_text, device):
    dataset = TextDataset(sample_text, sequence_length=3)
    model = BigramLanguageModel(dataset.vocab_size).to(device)
    context = torch.zeros((1, 1), dtype=torch.long, device=device)

    # Generate with different temperatures
    temps = [0.1, 0.5, 1.0, 2.0]
    generated_texts = []

    for temp in temps:
        indices = model.generate(context, max_new_tokens=50, temperature=temp)
        generated_texts.append(dataset.decode(indices[0].tolist()))

    # Check that all generated texts are valid
    for text in generated_texts:
        assert all(c in dataset.stoi for c in text)


import torch
import pytest
from src.data import TextDataset

def test_dataset_initialization(sample_text):
    dataset = TextDataset(sample_text, sequence_length=3)
    assert dataset.sequence_length == 3
    assert dataset.vocab_size == len(set(sample_text))
    assert len(dataset.stoi) == len(dataset.itos)

def test_dataset_encoding_decoding(dataset, sample_text):
    # Test if we can encode and decode back to the original text
    indices = [dataset.stoi[ch] for ch in sample_text[:dataset.sequence_length]]
    decoded = dataset.decode(indices)
    assert decoded == sample_text[:dataset.sequence_length]

def test_dataset_getitem(dataset):
    x, y = dataset[0]
    assert isinstance(x, torch.Tensor)
    assert isinstance(y, torch.Tensor)
    assert x.shape == (dataset.sequence_length,)
    assert y.shape == (dataset.sequence_length,)
    assert (y[:-1] == x[1:]).all()  # Check if y is shifted by one position

def test_dataset_length(dataset, sample_text):
    expected_length = len(sample_text) - dataset.sequence_length
    assert len(dataset) == expected_length

def test_invalid_sequence_length():
    with pytest.raises(ValueError):
        TextDataset("short", sequence_length=10)


import torch
import pytest
from src.data import TextDataset

def test_dataset_initialization(sample_text):
    dataset = TextDataset(sample_text, sequence_length=3)
    assert dataset.sequence_length == 3
    assert dataset.vocab_size == len(set(sample_text))
    assert len(dataset.stoi) == len(dataset.itos)

def test_dataset_encoding_decoding(dataset, sample_text):
    # Test if we can encode and decode back to the original text
    indices = [dataset.stoi[ch] for ch in sample_text[:dataset.sequence_length]]
    decoded = dataset.decode(indices)
    assert decoded == sample_text[:dataset.sequence_length]

def test_dataset_getitem(dataset):
    x, y = dataset[0]
    assert isinstance(x, torch.Tensor)
    assert isinstance(y, torch.Tensor)
    assert x.shape == (dataset.sequence_length,)
    assert y.shape == (dataset.sequence_length,)
    assert (y[:-1] == x[1:]).all()  # Check if y is shifted by one position

def test_dataset_length(dataset, sample_text):
    expected_length = len(sample_text) - dataset.sequence_length
    assert len(dataset) == expected_length

def test_invalid_sequence_length():


def test_dataset_initialization(sample_text):
    dataset = TextDataset(sample_text, sequence_length=3)
    assert dataset.sequence_length == 3
    assert dataset.vocab_size == len(set(sample_text))
    assert len(dataset.stoi) == len(dataset.itos)

def test_dataset_encoding_decoding(dataset, sample_text):
    # Test if we can encode and decode back to the original text
    indices = [dataset.stoi[ch] for ch in sample_text[:dataset.sequence_length]]
    decoded = dataset.decode(indices)
    assert decoded == sample_text[:dataset.sequence_length]

def test_dataset_getitem(dataset):
    x, y = dataset[0]
    assert isinstance(x, torch.Tensor)
    assert isinstance(y, torch.Tensor)
    assert x.shape == (dataset.sequence_length,)
    assert y.shape == (dataset.sequence_length,)
    assert (y[:-1] == x[1:]).all()  # Check if y is shifted by one position

def test_dataset_length(dataset, sample_text):
    expected_length = len(sample_text) - dataset.sequence_length
    assert len(dataset) == expected_length

def test_invalid_sequence_length():
    with pytest.raises(ValueError):
        TextDataset("short", sequence_length=10)


    assert len(dataset.stoi) == len(dataset.itos)

def test_dataset_encoding_decoding(dataset, sample_text):
    # Test if we can encode and decode back to the original text
    indices = [dataset.stoi[ch] for ch in sample_text[:dataset.sequence_length]]
    decoded = dataset.decode(indices)
    assert decoded == sample_text[:dataset.sequence_length]

def test_dataset_getitem(dataset):
    x, y = dataset[0]
    assert isinstance(x, torch.Tensor)
    assert isinstance(y, torch.Tensor)
    assert x.shape == (dataset.sequence_length,)
    assert y.shape == (dataset.sequence_length,)
    assert (y[:-1] == x[1:]).all()  # Check if y is shifted by one position

def test_dataset_length(dataset, sample_text):
    expected_length = len(sample_text) - dataset.sequence_length
    assert len(dataset) == expected_length

def test_invalid_sequence_length():
    with pytest.raises(ValueError):
        TextDataset("short", sequence_length=10)


src.utils.config

from dataclasses import dataclass
import torch

@dataclass
class TrainingConfig:
    input_file: str
    output_model: str = 'model.pt'
    sequence_length: int = 8
    batch_size: int = 32
    epochs: int = 50
    learning_rate: float = 0.001
    sample_every: int = 10
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'

@dataclass
class GenerationConfig:
    model_file: str
    input_file: str
    num_tokens: int = 100
    temperature: float = 1.0
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'


src.data.dataset

import torch
from torch.utils.data import Dataset
from collections import Counter
import numpy as np

class TextDataset(Dataset):
    def __init__(self, text, sequence_length=8):
        if not isinstance(sequence_length, int) or sequence_length <= 0:
            raise ValueError("sequence_length must be a positive integer")
        self.sequence_length = sequence_length

        # Create vocabulary
        if not isinstance(text, str) or len(text) == 0:
            raise ValueError("text must be a non-empty string")
        chars = sorted(list(set(text)))
        self.vocab_size = len(chars)
        self.stoi = {ch: i for i, ch in enumerate(chars)}
        self.itos = {i: ch for i, ch in enumerate(chars)}

        # Encode text
        data = torch.tensor([self.stoi[ch] for ch in text], dtype=torch.long)

        # Validate sequence_length relative to data length
        if len(data) <= sequence_length:
            raise ValueError("sequence_length must be less than the length of the text")

        # Create sequences
        self.sequences = []
        for i in range(0, len(data) - sequence_length):
            sequence = data[i:i + sequence_length + 1]
            self.sequences.append(sequence)

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        sequence = self.sequences[idx]
        x = sequence[:-1]
        y = sequence[1:]
        return x, y

    def decode(self, indices):
        return ''.join([self.itos[int(i)] for i in indices])


        if not isinstance(text, str) or len(text) == 0:
            raise ValueError("text must be a non-empty string")
        chars = sorted(list(set(text)))
        self.vocab_size = len(chars)
        self.stoi = {ch: i for i, ch in enumerate(chars)}
        self.itos = {i: ch for i, ch in enumerate(chars)}

        # Encode text
        data = torch.tensor([self.stoi[ch] for ch in text], dtype=torch.long)

        # Validate sequence_length relative to data length
        if len(data) <= sequence_length:
            raise ValueError("sequence_length must be less than the length of the text")

        # Create sequences
        self.sequences = []
        for i in range(0, len(data) - sequence_length):
            sequence = data[i:i + sequence_length + 1]
            self.sequences.append(sequence)

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        sequence = self.sequences[idx]
        x = sequence[:-1]
        y = sequence[1:]
        return x, y

    def decode(self, indices):
        return ''.join([self.itos[int(i)] for i in indices])


        self.vocab_size = len(chars)
        self.stoi = {ch: i for i, ch in enumerate(chars)}
        self.itos = {i: ch for i, ch in enumerate(chars)}

        # Encode text
        data = torch.tensor([self.stoi[ch] for ch in text], dtype=torch.long)

        # Validate sequence_length relative to data length
        if len(data) <= sequence_length:
            raise ValueError("sequence_length must be less than the length of the text")

        # Create sequences
        self.sequences = []
        for i in range(0, len(data) - sequence_length):
            sequence = data[i:i + sequence_length + 1]
            self.sequences.append(sequence)

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        sequence = self.sequences[idx]
        x = sequence[:-1]
        y = sequence[1:]
        return x, y

    def decode(self, indices):
        return ''.join([self.itos[int(i)] for i in indices])


decode

        # Validate sequence_length relative to data length
        if len(data) <= sequence_length:
            raise ValueError("sequence_length must be less than the length of the text")

        # Create sequences
        self.sequences = []
        for i in range(0, len(data) - sequence_length):
            sequence = data[i:i + sequence_length + 1]
            self.sequences.append(sequence)

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        sequence = self.sequences[idx]
        x = sequence[:-1]
        y = sequence[1:]
        return x, y

    def decode(self, indices):
        return ''.join([self.itos[int(i)] for i in indices])


import torch
import torch.nn as nn

class BigramLanguageModel(nn.Module):
    """A bigram language model implementation.

    A bigram language model is one of the simplest forms of language models that predicts
    the probability of a token based solely on the previous token. For a sequence of tokens
    [w₁, w₂, ..., wₙ], the bigram model approximates the probability as:

    P(w₁, w₂, ..., wₙ) ≈ P(w₁) * P(w₂|w₁) * P(w₃|w₂) * ... * P(wₙ|wₙ₋₁)

    The model maintains a transition matrix (bigram_table) where:
    - Each row i represents the current token
    - Each column j represents the next token
    - Entry (i,j) contains the logit for P(token_j|token_i)

    Args:
        vocab_size (int): Size of the vocabulary

    Attributes:
        bigram_table (nn.Parameter): Transition matrix of shape (vocab_size, vocab_size)
            containing the log-probabilities of each token following another token
        vocab_size (int): Size of the vocabulary
    """

    def __init__(self, vocab_size):
        super().__init__()
        # Initialize transition matrix with zeros
        # After training, bigram_table[i,j] will contain the log-probability
        # of token j following token i
        self.bigram_table = nn.Parameter(torch.zeros(vocab_size, vocab_size))
        self.vocab_size = vocab_size

    def forward(self, idx):
        """Compute the bigram probabilities for a batch of sequences.

        For each position in the input sequence, computes logits for the probability
        distribution over the next token.

        Args:
            idx (torch.Tensor): Batch of token sequences of shape (batch_size, sequence_length)

        Returns:
            torch.Tensor: Logits for next-token predictions of shape
                (batch_size, sequence_length, vocab_size)
        """
        B, T = idx.shape

        # Index into bigram_table to get next-token logits for each position
        # Ensure the parameter is on the same device as the indices
        table = self.bigram_table.to(idx.device)
        # Select rows for each token id efficiently and reshape
        rows = table.index_select(0, idx.view(-1))  # (B*T, vocab_size)
        logits = rows.view(B, T, self.vocab_size)

        return logits

    def generate(self, idx, max_new_tokens, temperature=1.0):
        """Generate new tokens autoregressively using the bigram probabilities.

        Starting from the last token in the provided context, generates new tokens
        one at a time by:
        1. Getting the probability distribution for the next token
        2. Applying temperature scaling to control randomness
        3. Sampling from the resulting distribution

        Args:
            idx (torch.Tensor): Starting context of shape (batch_size, sequence_length)
            max_new_tokens (int): Number of new tokens to generate
            temperature (float, optional): Temperature for sampling. Higher values make the
                distribution more uniform, lower values make it more peaky. Defaults to 1.0.

        Returns:
            torch.Tensor: Generated sequence including the context,
                shape (batch_size, sequence_length + max_new_tokens)
        """
        # Ensure the parameter is on the same device as the indices
        table = self.bigram_table.to(idx.device)

        for _ in range(max_new_tokens):
            # Only need the last token for bigram prediction
            last_token = idx[:, -1:]  # (B, 1)
            rows = table.index_select(0, last_token.view(-1))  # (B*1, vocab_size)
            logits = rows.view(idx.size(0), 1, self.vocab_size)  # (B, 1, vocab_size)

            # Apply temperature scaling to logits
            # Higher temperature = more random, Lower = more deterministic
            logits = logits[:, -1, :] / temperature  # (B, vocab_size)

            # Convert logits to probabilities
            probs = torch.softmax(logits, dim=-1)  # (B, vocab_size)

            # Sample next token from the probability distribution
            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)

            # Append new token to the sequence
            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)

        return idx


    A bigram language model is one of the simplest forms of language models that predicts
    the probability of a token based solely on the previous token. For a sequence of tokens
    [w₁, w₂, ..., wₙ], the bigram model approximates the probability as:

    P(w₁, w₂, ..., wₙ) ≈ P(w₁) * P(w₂|w₁) * P(w₃|w₂) * ... * P(wₙ|wₙ₋₁)

    The model maintains a transition matrix (bigram_table) where:
    - Each row i represents the current token
    - Each column j represents the next token
    - Entry (i,j) contains the logit for P(token_j|token_i)

    Args:
        vocab_size (int): Size of the vocabulary

    Attributes:
        bigram_table (nn.Parameter): Transition matrix of shape (vocab_size, vocab_size)
            containing the log-probabilities of each token following another token
        vocab_size (int): Size of the vocabulary
    """

    def __init__(self, vocab_size):
        super().__init__()
        # Initialize transition matrix with zeros
        # After training, bigram_table[i,j] will contain the log-probability
        # of token j following token i
        self.bigram_table = nn.Parameter(torch.zeros(vocab_size, vocab_size))
        self.vocab_size = vocab_size

    def forward(self, idx):
        """Compute the bigram probabilities for a batch of sequences.

        For each position in the input sequence, computes logits for the probability
        distribution over the next token.

        Args:
            idx (torch.Tensor): Batch of token sequences of shape (batch_size, sequence_length)

        Returns:
            torch.Tensor: Logits for next-token predictions of shape
                (batch_size, sequence_length, vocab_size)
        """
        B, T = idx.shape

        # Index into bigram_table to get next-token logits for each position
        # Ensure the parameter is on the same device as the indices
        table = self.bigram_table.to(idx.device)
        # Select rows for each token id efficiently and reshape


forward

    - Each column j represents the next token
    - Entry (i,j) contains the logit for P(token_j|token_i)

    Args:
        vocab_size (int): Size of the vocabulary

    Attributes:
        bigram_table (nn.Parameter): Transition matrix of shape (vocab_size, vocab_size)
            containing the log-probabilities of each token following another token
        vocab_size (int): Size of the vocabulary
    """

    def __init__(self, vocab_size):
        super().__init__()
        # Initialize transition matrix with zeros
        # After training, bigram_table[i,j] will contain the log-probability
        # of token j following token i
        self.bigram_table = nn.Parameter(torch.zeros(vocab_size, vocab_size))
        self.vocab_size = vocab_size

    def forward(self, idx):
        """Compute the bigram probabilities for a batch of sequences.

        For each position in the input sequence, computes logits for the probability
        distribution over the next token.

        Args:
            idx (torch.Tensor): Batch of token sequences of shape (batch_size, sequence_length)

        Returns:
            torch.Tensor: Logits for next-token predictions of shape
                (batch_size, sequence_length, vocab_size)
        """
        B, T = idx.shape

        # Index into bigram_table to get next-token logits for each position
        # Ensure the parameter is on the same device as the indices
        table = self.bigram_table.to(idx.device)
        # Select rows for each token id efficiently and reshape
        rows = table.index_select(0, idx.view(-1))  # (B*T, vocab_size)
        logits = rows.view(B, T, self.vocab_size)

        return logits

    def generate(self, idx, max_new_tokens, temperature=1.0):
        """Generate new tokens autoregressively using the bigram probabilities.

        Starting from the last token in the provided context, generates new tokens
        one at a time by:
        1. Getting the probability distribution for the next token
        2. Applying temperature scaling to control randomness
        3. Sampling from the resulting distribution

        Args:
            idx (torch.Tensor): Starting context of shape (batch_size, sequence_length)
            max_new_tokens (int): Number of new tokens to generate
            temperature (float, optional): Temperature for sampling. Higher values make the
                distribution more uniform, lower values make it more peaky. Defaults to 1.0.

        Returns:
            torch.Tensor: Generated sequence including the context,
                shape (batch_size, sequence_length + max_new_tokens)
        """


generate

        distribution over the next token.

        Args:
            idx (torch.Tensor): Batch of token sequences of shape (batch_size, sequence_length)

        Returns:
            torch.Tensor: Logits for next-token predictions of shape
                (batch_size, sequence_length, vocab_size)
        """
        B, T = idx.shape

        # Index into bigram_table to get next-token logits for each position
        # Ensure the parameter is on the same device as the indices
        table = self.bigram_table.to(idx.device)
        # Select rows for each token id efficiently and reshape
        rows = table.index_select(0, idx.view(-1))  # (B*T, vocab_size)
        logits = rows.view(B, T, self.vocab_size)

        return logits

    def generate(self, idx, max_new_tokens, temperature=1.0):
        """Generate new tokens autoregressively using the bigram probabilities.

        Starting from the last token in the provided context, generates new tokens
        one at a time by:
        1. Getting the probability distribution for the next token
        2. Applying temperature scaling to control randomness
        3. Sampling from the resulting distribution

        Args:
            idx (torch.Tensor): Starting context of shape (batch_size, sequence_length)
            max_new_tokens (int): Number of new tokens to generate
            temperature (float, optional): Temperature for sampling. Higher values make the
                distribution more uniform, lower values make it more peaky. Defaults to 1.0.

        Returns:
            torch.Tensor: Generated sequence including the context,
                shape (batch_size, sequence_length + max_new_tokens)
        """
        # Ensure the parameter is on the same device as the indices
        table = self.bigram_table.to(idx.device)

        for _ in range(max_new_tokens):
            # Only need the last token for bigram prediction
            last_token = idx[:, -1:]  # (B, 1)
            rows = table.index_select(0, last_token.view(-1))  # (B*1, vocab_size)
            logits = rows.view(idx.size(0), 1, self.vocab_size)  # (B, 1, vocab_size)

            # Apply temperature scaling to logits
            # Higher temperature = more random, Lower = more deterministic
            logits = logits[:, -1, :] / temperature  # (B, vocab_size)

            # Convert logits to probabilities
            probs = torch.softmax(logits, dim=-1)  # (B, vocab_size)

            # Sample next token from the probability distribution
            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)

            # Append new token to the sequence
            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)

        return idx


scripts.generate

from src.model import BigramLanguageModel
from src.data import TextDataset
from src.utils.config import GenerationConfig


scripts.train

from src.model import BigramLanguageModel
from src.data import TextDataset
from src.utils.config import TrainingConfig


setup.py

from setuptools import setup, find_packages

setup(
    name="bigram-lm",
    version="0.1.0",
    packages=find_packages(),
    install_requires=[
        "torch>=2.0.0",
        "numpy>=1.21.0",
    ],
    extras_require={
        'test': [
            'pytest>=7.0.0',
            'pytest-cov>=4.0.0',
        ],
    },
    author="Your Name",
    author_email="your.email@example.com",
    description="A simple bigram language model implementation",
    long_description=open("README.md").read(),
    long_description_content_type="text/markdown",


tests/conftest.py

import pytest
import torch
from src.model import BigramLanguageModel
from src.data import TextDataset

@pytest.fixture
def sample_text():
    return "Hello, World!"

@pytest.fixture
def dataset(sample_text):
    return TextDataset(sample_text, sequence_length=3)

@pytest.fixture
def model(dataset):
    return BigramLanguageModel(dataset.vocab_size)

@pytest.fixture
def device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'


import torch
import pytest
from src.model import BigramLanguageModel

def test_model_initialization(model, dataset):
    assert isinstance(model, BigramLanguageModel)
    assert model.bigram_table.shape == (dataset.vocab_size, dataset.vocab_size)

def test_model_forward(model, device):
    batch_size, seq_length = 2, 4
    idx = torch.randint(0, model.vocab_size, (batch_size, seq_length)).to(device)
    output = model(idx)

    assert output.shape == (batch_size, seq_length, model.vocab_size)
    assert not torch.isnan(output).any()

def test_model_generate(model, device):
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    max_new_tokens = 10

    generated = model.generate(context, max_new_tokens)


import torch
from src.model import BigramLanguageModel
from src.data import TextDataset

def test_end_to_end_generation(sample_text, device):
    # Create dataset and model
    dataset = TextDataset(sample_text, sequence_length=3)
    model = BigramLanguageModel(dataset.vocab_size).to(device)

    # Generate text
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    generated_indices = model.generate(context, max_new_tokens=20)
    generated_text = dataset.decode(generated_indices[0].tolist())

    assert len(generated_text) == 21  # context + max_new_tokens
    assert all(c in dataset.stoi for c in generated_text)

def test_generation_with_different_temperatures(sample_text, device):
    dataset = TextDataset(sample_text, sequence_length=3)
    model = BigramLanguageModel(dataset.vocab_size).to(device)
    context = torch.zeros((1, 1), dtype=torch.long, device=device)


import torch
import pytest
from src.data import TextDataset

def test_dataset_initialization(sample_text):
    dataset = TextDataset(sample_text, sequence_length=3)
    assert dataset.sequence_length == 3
    assert dataset.vocab_size == len(set(sample_text))
    assert len(dataset.stoi) == len(dataset.itos)

def test_dataset_encoding_decoding(dataset, sample_text):
    # Test if we can encode and decode back to the original text
    indices = [dataset.stoi[ch] for ch in sample_text[:dataset.sequence_length]]
    decoded = dataset.decode(indices)
    assert decoded == sample_text[:dataset.sequence_length]

def test_dataset_getitem(dataset):
    x, y = dataset[0]
    assert isinstance(x, torch.Tensor)
    assert isinstance(y, torch.Tensor)
    assert x.shape == (dataset.sequence_length,)


# Empty file to make the directory a Python package


from dataclasses import dataclass
import torch

@dataclass
class TrainingConfig:
    input_file: str
    output_model: str = 'model.pt'
    sequence_length: int = 8
    batch_size: int = 32
    epochs: int = 50
    learning_rate: float = 0.001
    sample_every: int = 10
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'

@dataclass
class GenerationConfig:
    model_file: str
    input_file: str
    num_tokens: int = 100
    temperature: float = 1.0
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'


import torch
from torch.utils.data import Dataset
from collections import Counter
import numpy as np

class TextDataset(Dataset):
    def __init__(self, text, sequence_length=8):
        if not isinstance(sequence_length, int) or sequence_length <= 0:
            raise ValueError("sequence_length must be a positive integer")
        self.sequence_length = sequence_length

        # Create vocabulary
        if not isinstance(text, str) or len(text) == 0:
            raise ValueError("text must be a non-empty string")
        chars = sorted(list(set(text)))
        self.vocab_size = len(chars)
        self.stoi = {ch: i for i, ch in enumerate(chars)}
        self.itos = {i: ch for i, ch in enumerate(chars)}

        # Encode text
        data = torch.tensor([self.stoi[ch] for ch in text], dtype=torch.long)


from .dataset import TextDataset

__all__ = ['TextDataset']


import torch
import torch.nn as nn

class BigramLanguageModel(nn.Module):
    """A bigram language model implementation.

    A bigram language model is one of the simplest forms of language models that predicts
    the probability of a token based solely on the previous token. For a sequence of tokens
    [w₁, w₂, ..., wₙ], the bigram model approximates the probability as:

    P(w₁, w₂, ..., wₙ) ≈ P(w₁) * P(w₂|w₁) * P(w₃|w₂) * ... * P(wₙ|wₙ₋₁)

    The model maintains a transition matrix (bigram_table) where:
    - Each row i represents the current token
    - Each column j represents the next token
    - Entry (i,j) contains the logit for P(token_j|token_i)

    Args:
        vocab_size (int): Size of the vocabulary

    Attributes:


from .bigram_model import BigramLanguageModel

__all__ = ['BigramLanguageModel']