import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from model import ResumeEncoder
from tokenizer import encode

# Load resume text
with open("resume_text.txt", "r", encoding="utf-8") as f:
    text = f.read()

data = encode(text)
context_length = 64

class ResumeDataset(Dataset):
    def __init__(self, data, context_length):
        self.data = data
        self.context_length = context_length

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

    def __getitem__(self, idx):
        x = torch.tensor(self.data[idx : idx + self.context_length])
        y = torch.tensor(self.data[idx + 1 : idx + self.context_length + 1])
        return x, y

dataset    = ResumeDataset(data, context_length)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

model     = ResumeEncoder()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
epochs    = 500


for epoch in range(1, epochs + 1):
    total_loss = 0
    for x, y in dataloader:
        logits = model(x)
        B, T, C = logits.shape
        loss = F.cross_entropy(logits.view(B*T, C), y.view(B*T))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    avg_loss = total_loss / len(dataloader)

    if epoch % 10 == 0:
        print(f"Epoch {epoch}/{epochs}  |  Loss: {avg_loss:.4f}")

    if epoch % 50 == 0:
        torch.save(model.state_dict(), f"resume_encoder_epoch{epoch}.pth")
        print(f"  → Checkpoint saved at epoch {epoch}")

torch.save(model.state_dict(), "resume_encoder.pth")
print("Final model saved!")