vibhu-llm

vibhu-llm is a PyTorch-based library providing a from-scratch implementation of a Large Language Model (LLM). The primary goal of this library is to offer a clear and understandable decoder-only transformer model that can be trained on custom datasets and used for text generation.

Installation

You can install vibhu-llm directly from PyPI:

pip install vibhu-llm

Usage

Here's a quick example of how to use vibhu-llm to train a model and generate text:

import torch
from llm.config import Config
from llm.data import get_train_and_val_dataloaders, get_tokenizer, DATA_FETCHERS
from llm.model import LLMModel
from llm.model_run import get_optimizer, train_model
from llm.utils import get_device
from llm.inference import generate_and_print_sample

# 1. Setup device and tokenizer
device = get_device()
print(f"Using device: {device}")
tokenizer = get_tokenizer()

# 2. Configure the model
config = Config(
    context_length=256,
    vocab_size=tokenizer.n_vocab,
    emb_dim=128,
    n_blocks=2,
    n_heads=8,
    dropout=0.1,
)
model = LLMModel(config).to(device)
print(f"Total model parameters: {sum(p.numel() for p in model.parameters()):,}")

# 3. Load data
# You can choose from available data fetchers
train_loader, val_loader = get_train_and_val_dataloaders(
    DATA_FETCHERS["harry_potter"],
    tokenizer,
    batch_size=32,
    max_length=config.context_length,
    stride=128,
    train_ratio=0.9,
)

# 4. Setup optimizer
optimizer = get_optimizer(model, lr=0.001)

# 5. Train the model
train_model(
    train_loader,
    val_loader,
    model,
    optimizer,
    num_epochs=10,
    config_dict=config.dict(),
    experiment_name="runs/my_experiment",
    eval_interval=100,
    save_interval=100,
)

# 6. Generate text
generate_and_print_sample(
    model,
    tokenizer,
    device,
    start_context="The magic of this world",
    context_size=config.context_length,
    max_new_tokens=50,
)

Configuration

The llm.config.Config class allows you to customize the model architecture. Here are the available parameters:

Contributing

If you're interested in improving the library, here's how you can get started:

1. Clone the repository:

   git clone https://github.com/Vibhu-Agarwal/LLM
   cd LLM
   

2. Install dependencies: It is recommended to use Poetry for managing dependencies.

   poetry install
   

3. Project Structure: