File size: 4,432 Bytes
db5e0ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import time
import glob
from tqdm import tqdm
from config import *
from model import PegeModel
from tokenizer import TurkishTokenizer
from dataset import TextDataset
def train():
print("=" * 50)
print("PEGE TRAINING STARTED")
print("=" * 50)
# Tokenizer
tokenizer = TurkishTokenizer(vocab_size=VOCAB_SIZE)
tokenizer_path = os.path.expanduser(f"{MODEL_DIR}/tokenizer.pkl")
data_dir = os.path.expanduser(DATA_DIR)
txt_files = glob.glob(f"{data_dir}/**/*.txt", recursive=True)
if not os.path.exists(data_dir) or not txt_files:
print(f"\nUYARI: {data_dir} klasöründe .txt dosyası bulunamadı!")
print("Lütfen makalelerini TXT olarak buraya at:")
print(f" {data_dir}/")
return
tokenizer_json = tokenizer_path.replace(".pkl", ".json")
if os.path.exists(tokenizer_json):
tokenizer.load(tokenizer_path)
else:
tokenizer.train(txt_files)
tokenizer.save(tokenizer_path)
dataset = TextDataset(data_dir, tokenizer, max_length=MAX_SEQ_LEN)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
vocab_size = len(tokenizer.stoi)
print(f"Vocabulary size: {vocab_size}")
# Model
model = PegeModel(vocab_size).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
# Checkpoint yükle
start_step = 0
checkpoint_files = glob.glob(os.path.expanduser(f"{CHECKPOINT_DIR}/checkpoint_*.pt"))
if checkpoint_files:
latest = max(checkpoint_files, key=os.path.getctime)
checkpoint = torch.load(latest, map_location=device)
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
start_step = checkpoint['step']
print(f"Resumed from checkpoint: {latest} (step {start_step})")
model.train()
# Training loop
step = start_step
losses = []
print(f"\nTraining on {len(dataset)} samples")
print(f"Batch size: {BATCH_SIZE}")
print(f"Device: {device}\n")
progress = tqdm(total=MAX_STEPS, initial=start_step, desc="Training")
while step < MAX_STEPS:
for batch_x, batch_y in dataloader:
batch_x = batch_x.to(device)
batch_y = batch_y.to(device)
# Forward
logits, loss = model(batch_x, batch_y)
# Backward
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
losses.append(loss.item())
# Logging
if step % 100 == 0:
avg_loss = sum(losses[-100:]) / min(len(losses), 100)
progress.set_postfix({'loss': f'{avg_loss:.4f}'})
# Checkpoint
if step % CHECKPOINT_INTERVAL == 0 and step > 0:
os.makedirs(os.path.expanduser(CHECKPOINT_DIR), exist_ok=True)
checkpoint_path = os.path.expanduser(f"{CHECKPOINT_DIR}/checkpoint_{step}.pt")
torch.save({
'step': step,
'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict(),
'vocab_size': vocab_size,
}, checkpoint_path)
# Eski checkpoint'leri sil
old_checkpoints = sorted(glob.glob(os.path.expanduser(f"{CHECKPOINT_DIR}/checkpoint_*.pt")))
for old in old_checkpoints[:-3]: # Son 3'ü tut
os.remove(old)
print(f"\nCheckpoint saved: {checkpoint_path}")
step += 1
progress.update(1)
if step >= MAX_STEPS:
break
progress.close()
# Final save
os.makedirs(os.path.expanduser(MODEL_DIR), exist_ok=True)
final_path = os.path.expanduser(f"{MODEL_DIR}/pege_final.pt")
torch.save({
'model_state': model.state_dict(),
'vocab_size': vocab_size,
'config': {
'embed_dim': EMBED_DIM,
'num_heads': NUM_HEADS,
'num_layers': NUM_LAYERS,
'hidden_dim': HIDDEN_DIM,
}
}, final_path)
print(f"\n{'=' * 50}")
print(f"Training complete! Model saved to: {final_path}")
print(f"{'=' * 50}")
if __name__ == "__main__":
train() |