| import torch |
| import sentencepiece as spm |
| from gpt_model import MiniGPT |
| import math |
| import time |
| import os |
|
|
| |
| |
| |
| SEQ_LEN = 128 |
| BATCH_SIZE = 32 |
| EPOCHS = 4 |
| LR = 3e-4 |
| MAX_TOKENS_LIMIT = 3000000 |
| CHECKPOINT_FILE = "/content/drive/MyDrive/xr_ai_checkpoint.pt" |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| print("Using device:", device) |
|
|
| |
| |
| |
| print("📦 Loading tokenizer...") |
|
|
| sp = spm.SentencePieceProcessor() |
| sp.load("hindi_tokenizer.model") |
|
|
| vocab_size = sp.get_piece_size() |
|
|
| |
| |
| |
| print("📂 Loading dataset...") |
|
|
| text = open("sangraha_hindi_final_dataset.txt", encoding="utf-8").read() |
|
|
| tokens = sp.encode(text) |
|
|
| if MAX_TOKENS_LIMIT: |
| tokens = tokens[:MAX_TOKENS_LIMIT] |
|
|
| data = torch.tensor(tokens, dtype=torch.long) |
|
|
| split = int(0.9 * len(data)) |
|
|
| train_data = data[:split] |
| val_data = data[split:] |
|
|
| |
| |
| |
| model = MiniGPT(vocab_size).to(device) |
|
|
| optimizer = torch.optim.AdamW(model.parameters(), lr=LR) |
|
|
| loss_fn = torch.nn.CrossEntropyLoss() |
|
|
| resume_step = 0 |
|
|
| |
| |
| |
| if os.path.exists(CHECKPOINT_FILE): |
|
|
| print("🔄 Loading checkpoint...") |
|
|
| checkpoint = torch.load(CHECKPOINT_FILE) |
|
|
| model.load_state_dict(checkpoint["model"]) |
| optimizer.load_state_dict(checkpoint["optimizer"]) |
|
|
| resume_step = checkpoint["step"] |
|
|
| print("Resuming from step:", resume_step) |
|
|
| |
| |
| |
| print("🚀 XR_AI Training Started...") |
|
|
| step = resume_step |
|
|
| for epoch in range(EPOCHS): |
|
|
| model.train() |
|
|
| total_train_loss = 0 |
| train_steps = 0 |
|
|
| for i in range(0, len(train_data) - SEQ_LEN - 1, SEQ_LEN * BATCH_SIZE): |
|
|
| batch_inputs = [] |
| batch_targets = [] |
|
|
| for j in range(BATCH_SIZE): |
|
|
| start = i + j * SEQ_LEN |
|
|
| if start + SEQ_LEN + 1 >= len(train_data): |
| continue |
|
|
| x = train_data[start:start+SEQ_LEN] |
| y = train_data[start+1:start+SEQ_LEN+1] |
|
|
| batch_inputs.append(x) |
| batch_targets.append(y) |
|
|
| if len(batch_inputs) == 0: |
| continue |
|
|
| inputs = torch.stack(batch_inputs).to(device) |
| targets = torch.stack(batch_targets).to(device) |
|
|
| optimizer.zero_grad() |
|
|
| outputs = model(inputs) |
|
|
| loss = loss_fn(outputs.view(-1, vocab_size), targets.view(-1)) |
|
|
| loss.backward() |
|
|
| optimizer.step() |
|
|
| step += 1 |
| train_steps += 1 |
| total_train_loss += loss.item() |
|
|
| |
| if step % 200 == 0: |
|
|
| torch.save({ |
| "model": model.state_dict(), |
| "optimizer": optimizer.state_dict(), |
| "step": step |
| }, CHECKPOINT_FILE) |
|
|
| print("💾 checkpoint saved at step", step) |
|
|
| avg_train_loss = total_train_loss / train_steps |
|
|
| |
| |
| |
| model.eval() |
|
|
| total_val_loss = 0 |
| val_steps = 0 |
|
|
| with torch.no_grad(): |
|
|
| for i in range(0, len(val_data) - SEQ_LEN - 1, SEQ_LEN): |
|
|
| inputs = val_data[i:i+SEQ_LEN].unsqueeze(0).to(device) |
| targets = val_data[i+1:i+SEQ_LEN+1].unsqueeze(0).to(device) |
|
|
| outputs = model(inputs) |
|
|
| loss = loss_fn(outputs.view(-1, vocab_size), targets.view(-1)) |
|
|
| total_val_loss += loss.item() |
| val_steps += 1 |
|
|
| avg_val_loss = total_val_loss / val_steps |
|
|
| perplexity = math.exp(avg_val_loss) |
|
|
| print( |
| f"Epoch {epoch+1}/{EPOCHS} | " |
| f"Train: {avg_train_loss:.4f} | " |
| f"Val: {avg_val_loss:.4f} | " |
| f"PPL: {perplexity:.2f}" |
| ) |
|
|
| torch.save(model.state_dict(), "best_model.pt") |
|
|
| print("🏁 Training complete") |
|
|
