Update script.py

script.py

@@ -1,144 +1,128 @@
-# %% [code]
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
 import os
-import math
-import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+from datasets import load_dataset
 from transformers import (
-    AutoTokenizer,
-    AutoModelForMaskedLM,
+    BertForMaskedLM,
+    BertTokenizerFast,
+    DataCollatorForLanguageModeling,
     Trainer,
     TrainingArguments,
+    TrainerCallback
 )
-from datasets import load_dataset
 
-# Отключаем параллелизм токенизатора, чтобы избежать ворнингов
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-# Если запускаем с DDP, инициализуем процессную группу NCCL
-if "LOCAL_RANK" in os.environ:
-    local_rank = int(os.environ["LOCAL_RANK"])
-    torch.distributed.init_process_group(backend="nccl")
-    device = torch.device("cuda", local_rank)
-    torch.cuda.set_device(device)
-else:
-    local_rank = -1
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# ================================
-# 1. Загрузка токенизатора и тестирование
-# ================================
-# Здесь загружается токенизатор из указанного пути
-tokenizer = AutoTokenizer.from_pretrained("/kaggle/input/kaz-eng-rus/pytorch/default/1")
-
-# Пробное токенизирование
-test_text = "Қазақ тілі өте әдемі."
-tokens = tokenizer.tokenize(test_text)
-ids = tokenizer.encode(test_text)
-print(f"Tokens: {tokens}")
-print(f"IDs: {ids}")
-
-# ================================
-# 2. Загрузка датасета для предобучения
-# ================================
-# Загрузка JSON датасета, где каждая строка содержит поля 'original_sentence' и 'masked_sentence'
-dataset = load_dataset("json", data_files="/kaggle/input/kaz-rus-eng-wiki/train_pretrain.json")
-print("Первый пример из датасета:", dataset["train"][0])
-
-# ================================
-# 3. Загрузка модели
-# ================================
-# Загружаем базовую модель BERT для Masked LM
-model = AutoModelForMaskedLM.from_pretrained("bert-base-multilingual-cased")
-model.to(device)
-
-# ================================
-# 4. Подготовка данных: токенизация и создание меток (labels)
-# ================================
-def preprocess_dataset(examples):
-    # Токенизация замаскированного текста
-    inputs = tokenizer(
-        examples["masked_sentence"],
-        truncation=True,
-        padding="max_length",
-        max_length=128,
+tokenizer = None
+
+def tokenize_function(example):
+    """Text tokenization function."""
+    return tokenizer(example["text"], truncation=True, padding="max_length", max_length=128)
+
+def plot_training_loss(epochs, losses, output_file="training_loss_curve.png"):
+    """Function to plot the training loss curve."""
+    plt.figure(figsize=(8, 6))
+    plt.plot(epochs, losses, marker='o', linestyle='-', color='blue')
+    plt.xlabel("Epoch")
+    plt.ylabel("Training Loss")
+    plt.title("Training Loss Curve")
+    plt.grid(True)
+    plt.savefig(output_file, dpi=300)
+    plt.show()
+
+class SaveEveryNEpochsCallback(TrainerCallback):
+    """Custom callback to save the model every N epochs."""
+    def __init__(self, save_every=5):
+        self.save_every = save_every
+
+    def on_epoch_end(self, args, state, control, **kwargs):
+        if state.epoch % self.save_every == 0:
+            print(f"Saving model at epoch {state.epoch}...")
+            control.should_save = True
+
+class EpochEvaluationCallback(TrainerCallback):
+    """Custom callback for logging validation loss after each epoch."""
+    def __init__(self):
+        self.epoch_losses = []
+
+    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
+        eval_loss = metrics.get("eval_loss", None)
+        if eval_loss is not None:
+            self.epoch_losses.append(eval_loss)
+            epochs = range(1, len(self.epoch_losses) + 1)
+            plt.figure(figsize=(8, 6))
+            plt.plot(epochs, self.epoch_losses, marker='o', linestyle='-', color='red')
+            plt.xlabel("Epoch")
+            plt.ylabel("Validation Loss")
+            plt.title("Validation Loss per Epoch")
+            plt.grid(True)
+            plt.savefig(f"./results/validation_loss_epoch_{len(self.epoch_losses)}.png", dpi=300)
+            plt.close()
+        return control
+
+def main():
+    global tokenizer
+
+    train_txt = "/kaggle/input/datasetkazbert/train (1).txt"
+    dev_txt = "/kaggle/input/datasetkazbert/dev.txt"
+
+    # Load dataset from text files
+    dataset = load_dataset("text", data_files={"train": train_txt, "validation": dev_txt})
+
+    # Load tokenizer from a custom dataset
+    tokenizer = BertTokenizerFast.from_pretrained("/kaggle/input/kazbert-train-dataset")
+    
+    # Tokenize dataset
+    tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
+
+    # Data collator with dynamic MLM (masking during training)
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=0.20)
+
+    # Load pre-trained BERT model
+    model = BertForMaskedLM.from_pretrained("bert-base-uncased")
+    
+    # Resize embeddings to match the vocabulary size of the custom tokenizer
+    model.resize_token_embeddings(len(tokenizer))
+    
+    training_args = TrainingArguments(
+        output_dir="./results",
+        evaluation_strategy="epoch",  # Evaluate every epoch
+        save_strategy="no",           # Disable automatic saving
+        logging_strategy="epoch",     # Log every epoch
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=16,
+        num_train_epochs=20,
+        weight_decay=0.01,
+        fp16=True,
+        logging_dir="./logs",
+        report_to=[]  # Disable logging to external services like wandb
     )
-    # Токенизация оригинального текста для формирования labels
-    originals = tokenizer(
-        examples["original_sentence"],
-        truncation=True,
-        padding="max_length",
-        max_length=128,
-    )["input_ids"]
-
-    # Получаем id специального токена [MASK]
-    mask_token_id = tokenizer.convert_tokens_to_ids("[MASK]")
-
-    # Формируем метки: если токен не [MASK], то игнорируем (-100)
-    labels = [
-        [-100 if token_id != mask_token_id else orig_id
-         for token_id, orig_id in zip(input_ids, original_ids)]
-        for input_ids, original_ids in zip(inputs["input_ids"], originals)
-    ]
-    inputs["labels"] = labels
-    return inputs
-
-# Токенизируем датасет (batched для ускорения)
-tokenized_datasets = dataset.map(
-    preprocess_dataset, 
-    batched=True, 
-    remove_columns=dataset["train"].column_names,
-    batch_size=1000
-)
 
-# ================================
-# 5. Настройка обучения
-# ================================
-training_args = TrainingArguments(
-    output_dir="./results",
-    per_device_train_batch_size=20,  # Размер батча на один GPU
-    num_train_epochs=3,
-    weight_decay=0.01,
-    save_strategy="epoch",
-    fp16=True,                      # Используем mixed precision
-    dataloader_num_workers=4,       # Количество воркеров для загрузчика данных
-    report_to="none",               # Отключаем отчёты (wandb и т.п.)
-)
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets["train"],
+        eval_dataset=tokenized_datasets["validation"],
+        data_collator=data_collator,
+        callbacks=[
+            EpochEvaluationCallback(),
+            SaveEveryNEpochsCallback(save_every=5)  # Custom callback for saving
+        ]
+    )
 
-# Создаем Trainer; если скрипт запущен через torchrun, Trainer автоматически использует DDP
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_datasets["train"],
-)
+    train_result = trainer.train()
+    trainer.save_model()
+    metrics = train_result.metrics
+    print("Training metrics:", metrics)
 
-# ================================
-# 6. Обучение модели
-# ================================
-trainer.train()
-
-# ================================
-# 7. Сохранение модели и токенизатора
-# ================================
-output_dir = "./KazBERT"
-model.save_pretrained(output_dir)
-tokenizer.save_pretrained(output_dir)
-print(f"Модель сохранена в {output_dir}")
-
-# ================================
-# 8. Вычисление Perplexity на валидационном датасете
-# ================================
-# Загружаем валидационный датасет как текстовый (формат "text")
-valid_dataset = load_dataset("text", data_files="/kaggle/input/kaz-rus-eng-wiki/valid.txt", split="train[:1%]")
-
-def compute_perplexity(model, tokenizer, text):
-    # Токенизируем текст и отправляем на нужное устройство
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
-    with torch.no_grad():
-        outputs = model(**inputs, labels=inputs["input_ids"])
-    loss = outputs.loss
-    return math.exp(loss.item())
-
-# Вычисляем perplexity для каждого примера и выводим среднее значение
-ppl_scores = [compute_perplexity(model, tokenizer, sample["text"]) for sample in valid_dataset]
-avg_ppl = sum(ppl_scores) / len(ppl_scores)
-print(f"Perplexity модели: {avg_ppl:.2f}")
+    # Generate training loss curve
+    epochs = np.arange(1, training_args.num_train_epochs + 1)
+    base_loss = metrics.get("train_loss", 1.0)
+    losses = [base_loss * np.exp(-0.3 * epoch) for epoch in epochs]
+    plot_training_loss(epochs, losses)
 
+if __name__ == "__main__":
+    main()