Add training script for fine-tuned first KoBART

first_kobart/train_stt2pron_eos.py

@@ -0,0 +1,98 @@
+import os
+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import BartForConditionalGeneration, PreTrainedTokenizerFast
+from torch.optim import AdamW
+from transformers import get_scheduler
+from tqdm import tqdm
+
+# ✅ 설정
+MODEL_DIR = "gogamza/kobart-base-v2"
+SAVE_DIR = "./kobart_stt2pron_with_eos"
+DATA_PATH = "data/train_stt2pron_with_eos.pt"
+BATCH_SIZE = 8
+EPOCHS = 7
+LR = 5e-5
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# ✅ 데이터셋 클래스
+class STT2PronDataset(Dataset):
+    def __init__(self, data, tokenizer, max_length=128):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        source = item["stt"]
+        target = item["pronunciation"]
+        input_enc = self.tokenizer(
+            source,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt"
+        )
+        target_enc = self.tokenizer(
+            target,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt"
+        )
+        labels = target_enc["input_ids"]
+        labels[labels == self.tokenizer.pad_token_id] = -100  # CrossEntropy loss 무시
+
+        return {
+            "input_ids": input_enc["input_ids"].squeeze(),
+            "attention_mask": input_enc["attention_mask"].squeeze(),
+            "labels": labels.squeeze()
+        }
+
+# ✅ 모델 및 토크나이저 로드
+tokenizer = PreTrainedTokenizerFast.from_pretrained(MODEL_DIR)
+model = BartForConditionalGeneration.from_pretrained(MODEL_DIR).to(DEVICE)
+
+# ✅ 데이터 로드 및 데이터로더 생성
+data = torch.load(DATA_PATH)
+dataset = STT2PronDataset(data, tokenizer)
+loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+# ✅ 옵티마이저 & 스케줄러
+optimizer = AdamW(model.parameters(), lr=LR)
+lr_scheduler = get_scheduler(
+    name="linear", optimizer=optimizer, num_warmup_steps=0,
+    num_training_steps=len(loader) * EPOCHS
+)
+
+# ✅ 학습 루프
+model.train()
+for epoch in range(EPOCHS):
+    print(f"\n🌟 Epoch {epoch+1}/{EPOCHS}")
+    loop = tqdm(loader)
+    total_loss = 0
+    for batch in loop:
+        for k in batch:
+            batch[k] = batch[k].to(DEVICE)
+        outputs = model(**batch)
+        loss = outputs.loss
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        lr_scheduler.step()
+
+        total_loss += loss.item()
+        loop.set_description(f"Loss: {loss.item():.4f}")
+    
+    avg_loss = total_loss / len(loader)
+    print(f"✅ Epoch {epoch+1} 평균 Loss: {avg_loss:.4f}")
+
+# ✅ 모델 저장
+os.makedirs(SAVE_DIR, exist_ok=True)
+model.save_pretrained(SAVE_DIR)
+tokenizer.save_pretrained(SAVE_DIR)
+print(f"\n📦 모델 저장 완료: {SAVE_DIR}")