code/data/pg19/prepare.py

import torch
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset
import tiktoken
from typing import List, Dict, Any, Optional


class PG19Dataset(Dataset):
    """ Dataset class for test-split of PG-19. 
    Downloads the dataset from huggingface `emozilla/pg19-test` and preprocesses it 
    into chunks of a specified sequence length and tokenizes it using gpt2 tokenizer. 
    """
    def __init__(
        self, 
        seq_len: int = 1024,
        streaming: bool = False,
        max_books: Optional[int] = None
    ):
        self.seq_len = seq_len
        self.tokenizer = tiktoken.get_encoding("gpt2")
        
        # Load the dataset
        self.dataset = load_dataset("emozilla/pg19-test", split="test", streaming=streaming)
        
        # If not streaming and max_books is specified, limit the dataset
        if not streaming and max_books is not None:
            self.dataset = self.dataset.select(range(min(max_books, len(self.dataset))))
        
        # Preprocess and chunk the texts
        self.chunks = []
        self._prepare_chunks(max_books if streaming else None)
        
        print(f"Dataset prepared with {len(self.chunks)} chunks of max length {seq_len}")
    
    def _prepare_chunks(self, max_books: Optional[int] = None):
        book_count = 0
        
        for example in self.dataset:
            if max_books and book_count >= max_books:
                break
                
            text = example['text']
            
            # Tokenize the entire book
            tokens = self.tokenizer.encode(text)
            
            # Split into chunks of seq_len
            for i in range(0, len(tokens), self.seq_len):
                chunk = tokens[i:i + self.seq_len]
                
                # Only keep chunks that are at least seq_len tokens
                # (you might want to adjust this based on your needs)
                if len(chunk) == self.seq_len:
                    self.chunks.append({
                        'input_ids': torch.tensor(chunk, dtype=torch.long),
                        'book_id': example.get('book_id', book_count),
                        'chunk_id': i // self.seq_len
                    })
            
            book_count += 1
            if book_count % 10 == 0:
                print(f"Processed {book_count} books, {len(self.chunks)} chunks so far...")
    
    def __len__(self) -> int:
        return len(self.chunks)
    
    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        return self.chunks[idx]


class PG19DataLoader:
    def __init__(
        self,
        seq_len: int = 1024,
        batch_size: int = 1,
        streaming: bool = False,
        max_books: Optional[int] = None,
        num_workers: int = 4,
        shuffle: bool = True,
        pin_memory: bool = True
    ):
        self.dataset = PG19Dataset(
            seq_len=seq_len,
            streaming=streaming,
            max_books=max_books
        )
        
        self.dataloader = DataLoader(
            self.dataset,
            batch_size=batch_size,
            shuffle=shuffle and not streaming,  # Can't shuffle streaming datasets
            num_workers=num_workers,
            pin_memory=pin_memory,
            collate_fn=self._collate_fn
        )
    
    def _collate_fn(self, batch: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
        # Stack input_ids
        input_ids = torch.stack([item['input_ids'] for item in batch])
        
        # Collect metadata
        book_ids = torch.tensor([item['book_id'] for item in batch], dtype=torch.long)
        chunk_ids = torch.tensor([item['chunk_id'] for item in batch], dtype=torch.long)
        
        return {
            'input_ids': input_ids,
            'book_ids': book_ids,
            'chunk_ids': chunk_ids
        }
    
    def __iter__(self):
        return iter(self.dataloader)
    
    def __len__(self):
        return len(self.dataloader)


# Example usage and testing
if __name__ == "__main__":
    # Example 1: Basic usage
    print("Creating PG-19 DataLoader...")
    
    # For testing, use a small number of books
    pg19_loader = PG19DataLoader(
        seq_len=10240,
        batch_size=1,
        max_books=100,
        num_workers=4,
        shuffle=True
    )
    
    print(f"DataLoader created with {len(pg19_loader)} batches")
    
    # Test loading a few batches
    print("\nTesting batch loading...")
    for i, batch in enumerate(pg19_loader):
        print(f"  Batch {i+1}:")
        print(f"  Input IDs shape: {batch['input_ids'].shape}")
        print(f"  Book IDs: {batch['book_ids'].tolist()}")
        print(f"  Chunk IDs: {batch['chunk_ids'].tolist()}")
        
        # Decode first few tokens to verify
        tokenizer = tiktoken.get_encoding("gpt2")
        first_sequence = batch['input_ids'][0][:100]
        decoded_text = tokenizer.decode(first_sequence.tolist())
        print(f"  Sample text: {decoded_text}...")
        
        if i >= 2:  # Only show first 3 batches
            break
    
    # print("\n" + "="*50)
    # print("Example 2: Streaming mode for large datasets")
    
    # # Example with streaming (useful for very large datasets)
    # streaming_loader = PG19DataLoader(
    #     seq_len=1024,
    #     batch_size=2,
    #     streaming=True,
    #     max_books=3,  # Still limit for demo
    #     shuffle=False  # Can't shuffle streaming datasets
    # )
    
    # print(f"Streaming DataLoader created")
    
    # # Test a few batches from streaming
    # for i, batch in enumerate(streaming_loader):
    #     print(f"Streaming Batch {i+1}: {batch['input_ids'].shape}")
    #     if i >= 1:  # Only show first 2 batches
    #         break