#!curl -L -O https://www.kaggle.com/api/v1/datasets/download/duckypolice/elapticai-1024
#!apt-get install unzip
#!unzip ./elapticai-1024
#!rm ./elapticai-1024import os
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
import sentencepiece as spm

# Configuration parameters
NUM_LAYERS = 12
D_MODEL = 360
NUM_HEADS = 12
DFF = 2048  # Feed-forward network dimension
VOCAB_SIZE = 59129
SEQ_LENGTH = 128
CHECKPOINT_DIR = './checkpoints'  # Directory to save checkpoints
MAX_CHECKPOINTS = 3
SAVE_INTERVAL = 1000
BATCH_SIZE = 32

# Create a subword tokenizer
if not os.path.isfile("mymodel.model"):
    spm.SentencePieceTrainer.train(input='1024.txt', model_prefix='mymodel', vocab_size=VOCAB_SIZE)
sp = spm.SentencePieceProcessor(model_file='mymodel.model')

# Define a minimal transformer model
class TransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, num_layers, dff):
        super(TransformerModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=num_heads, dim_feedforward=dff)
        self.transformer = nn.TransformerEncoder(self.transformer_layer, num_layers=num_layers)
        self.fc_out = nn.Linear(d_model, vocab_size)

    def forward(self, src):
        x = self.embedding(src) * (self.embedding.embedding_dim ** 0.5)
        x = self.transformer(x)
        x = self.fc_out(x)
        return x

# A generator to simulate data streaming
class TextDatasetIter:
    def __init__(self, file_path, tokenizer):
        self.file_path = file_path
        self.tokenizer = tokenizer
        with open(file_path, 'r', encoding='utf-8') as f:
            self.text = f.read()
        self.data = self.tokenizer.encode(self.text, out_type=int)

    def __len__(self):
        return len(self.data) - SEQ_LENGTH

    def __iter__(self):
        for i in range(len(self.data) - SEQ_LENGTH):
            yield (
                torch.tensor(self.data[i:i + SEQ_LENGTH], dtype=torch.long),
                torch.tensor(self.data[i + 1:i + SEQ_LENGTH + 1], dtype=torch.long)
            )

# Function to count parameters
def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

# Function to save checkpoint
def save_checkpoint(model, optimizer, epoch, idx):
    if not os.path.exists(CHECKPOINT_DIR):
        os.makedirs(CHECKPOINT_DIR)

    checkpoint_path = os.path.join(CHECKPOINT_DIR, f'checkpoint_epoch_{epoch}_iter_{idx}.pt')
    torch.save({
        'epoch': epoch,
        'iteration': idx,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }, checkpoint_path)

    # Remove old checkpoints if exceeding max checkpoints
    checkpoints = sorted(os.listdir(CHECKPOINT_DIR))
    if len(checkpoints) > MAX_CHECKPOINTS:
        os.remove(os.path.join(CHECKPOINT_DIR, checkpoints[0]))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TransformerModel(VOCAB_SIZE, D_MODEL, NUM_HEADS, NUM_LAYERS, DFF).to(device)

optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Update to use the tokenizer
dataset_iterator = TextDatasetIter("1024.txt", sp)

# Use tqdm.write instead of print
tqdm.write(f"Total trainable parameters in the model: {count_parameters(model)}")

# Training loop with tqdm
def train_model(model, dataset_iterator, optimizer, criterion, num_epochs=10):
    model.train()
    iter_len = len(dataset_iterator)

    for epoch in range(num_epochs):
        batch_loss = 0
        current_batch = []

        # Initialize tqdm progress bar
        start_time = time.time()
        progress_bar = tqdm(total=iter_len, desc=f"Epoch {epoch + 1}", unit="it", ncols=100)

        for idx, (input_seq, target_seq) in enumerate(dataset_iterator):
            current_batch.append((input_seq, target_seq))

            if len(current_batch) == BATCH_SIZE:  # Process batch
                batch_inputs, batch_targets = zip(*current_batch)
                batch_inputs = torch.stack(batch_inputs).to(device)
                batch_targets = torch.stack(batch_targets).to(device)

                optimizer.zero_grad()
                output = model(batch_inputs.permute(1, 0))
                loss = criterion(output.view(-1, VOCAB_SIZE), batch_targets.view(-1))
                loss.backward()
                optimizer.step()

                batch_loss += loss.item()
                current_batch = []  # Reset batch

                # Update progress bar
                progress_bar.set_postfix(
                    loss=f"{loss.item():.3f}",
                    used=time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time)),
                    est=time.strftime("%H:%M:%S", time.gmtime((time.time() - start_time) / (idx + 1) * (iter_len - idx)))
                )
                progress_bar.update(BATCH_SIZE)

                if idx % SAVE_INTERVAL <= BATCH_SIZE and idx > BATCH_SIZE:
                    save_checkpoint(model, optimizer, epoch, idx)
        
        progress_bar.close()
        tqdm.write(f"Epoch {epoch+1} Complete with Average Loss: {batch_loss / len(dataset_iterator):.4f}")

train_model(model, dataset_iterator, optimizer, criterion)


alter this to use fp16 and also explain what i would be able to use this model for