data.py

import torch


def load_data(text_path: str, train_ratio: float):
    # read data
    with open(text_path, 'r', encoding='utf-8') as f:
        text = f.read()

    # tokenize and encode data
    _, encode, _ = tokenize(text)
    encoded_text = encode(text)
    data = torch.tensor(encoded_text, dtype=torch.long)

    # split data
    train_cutoff = int(train_ratio * len(data))
    train_data = data[:train_cutoff]
    val_data = data[train_cutoff:]

    # return data splits
    return train_data, val_data
    

def tokenize(text: str):
    # create vocabulary
    vocab_list = sorted(set(text))
    
    # create encoder and decoder dictionaries
    encode_dict = {char: i for i, char in enumerate(vocab_list)}
    decode_dict = {i: char for i, char in enumerate(vocab_list)}
    
    # create encoder and decoder functions
    encode = lambda string: [encode_dict[char] for char in string]
    decode = lambda tokens: ''.join([decode_dict[token] for token in tokens])
    
    # return vocabulary, encoder, decoder
    return vocab_list, encode, decode


def sample_batch(train_data: torch.Tensor, val_data: torch.Tensor, split: str, max_context: int, batch_size: int):
    # select proper split
    source_data = train_data if split == "train" else val_data

    # generate random sample points
    sample_points = torch.randint(0, len(source_data) - max_context, (batch_size,)).tolist()

    # generate training and target data
    x = torch.stack([source_data[sample:sample+max_context] for sample in sample_points])
    y = torch.stack([source_data[sample+1:sample+max_context+1] for sample in sample_points])

    # return training and target data
    return x, y