File size: 2,903 Bytes
c1b8df7 | 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 | """
PACE Actor Full Training Loop
"""
import os
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, AutoModelForCausalLM
from torch.optim import AdamW
from tqdm import tqdm
from dataset import CoderActorDataset
def main():
# ==========================================
# 1. HARDWARE SETUP
# ==========================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# ==========================================
# 2. DATA PIPELINE
# ==========================================
model_name = "bigcode/tiny_starcoder_py"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
dataset = CoderActorDataset(
parquet_path="masteries/coding/data/raw/stack_python_157k.parquet",
tokenizer=tokenizer,
max_length=512,
)
loader = DataLoader(dataset, batch_size=4, shuffle=True)
# ==========================================
# 3. NEURAL NETWORK SETUP
# ==========================================
print("Loading 164M Parameter Actor Model to GPU...")
model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
optimizer = AdamW(model.parameters(), lr=5e-5)
# ==========================================
# 4. THE FULL TRAINING LOOP
# ==========================================
epochs = 1
save_dir = "masteries/coding/models/actor_v1"
print(f"Starting Full Training Loop for {epochs} Epoch(s)...")
model.train()
for epoch in range(epochs):
print(f"\n--- Epoch {epoch + 1}/{epochs} ---")
progress_bar = tqdm(loader, desc="Training")
total_loss = 0
for batch in progress_bar:
input_ids = batch["input_ids"].to(device)
labels = batch["labels"].to(device)
attention_mask = batch["attention_mask"].to(device)
outputs = model(
input_ids=input_ids, attention_mask=attention_mask, labels=labels
)
loss = outputs.loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
current_loss = loss.item()
total_loss += current_loss
progress_bar.set_postfix({"loss": f"{current_loss:.4f}"})
avg_loss = total_loss / len(loader)
print(f"\nEpoch {epoch + 1} completed. Average Loss: {avg_loss:.4f}")
# ==========================================
# 5. SAVE THE MODEL WEIGHTS
# ==========================================
print(f"Saving model checkpoint to {save_dir}...")
os.makedirs(save_dir, exist_ok=True)
model.save_pretrained(save_dir)
tokenizer.save_pretrained(save_dir)
print("Model saved successfully!")
if __name__ == "__main__":
main()
|