Create train.py

train.py

@@ -0,0 +1,106 @@
+import os
+import numpy as np
+import transformers
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForTokenClassification, 
+    TrainingArguments, 
+    Trainer,
+    DataCollatorForTokenClassification
+)
+from datasets import load_dataset
+import evaluate
+
+# 1. Константы и параметры
+MODEL_CHECKPOINT = "Eraly-ml/KazBERT"
+TASK = "ner"
+BATCH_SIZE = 32
+LEARNING_RATE = 5e-5
+EPOCHS = 8
+SAVE_PATH = "./kazbert_ner_finetuned"
+
+# 2. Загрузка данных и метрик
+datasets = load_dataset("issai/kaznerd", trust_remote_code=True)
+metric = evaluate.load("seqeval")
+label_list = datasets["train"].features[f"{TASK}_tags"].feature.names
+
+# 3. Токенизация с выравниванием меток
+tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT, trust_remote_code=True)
+
+def tokenize_and_align_labels(examples):
+    tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
+    labels = []
+    for i, label in enumerate(examples[f"{TASK}_tags"]):
+        word_ids = tokenized_inputs.word_ids(batch_index=i)
+        previous_word_idx = None
+        label_ids = []
+        for word_idx in word_ids:
+            if word_idx is None:
+                label_ids.append(-100) # Игнорируем спецтокены
+            elif word_idx != previous_word_idx:
+                label_ids.append(label[word_idx]) # Первая часть слова
+            else:
+                label_ids.append(-100) # Остальные части слова (игнорируем для оценки)
+            previous_word_idx = word_idx
+        labels.append(label_ids)
+    tokenized_inputs["labels"] = labels
+    return tokenized_inputs
+
+tokenized_datasets = datasets.map(tokenize_and_align_labels, batched=True)
+
+# 4. Модель с маппингом меток
+id2label = {i: label for i, label in enumerate(label_list)}
+label2id = {label: i for i, label in enumerate(label_list)}
+
+model = AutoModelForTokenClassification.from_pretrained(
+    MODEL_CHECKPOINT, 
+    num_labels=len(label_list),
+    id2label=id2label,
+    label2id=label2id,
+    trust_remote_code=True
+)
+
+# 5. Функция для вычисления метрик (Precision, Recall, F1)
+def compute_metrics(p):
+    predictions, labels = p
+    predictions = np.argmax(predictions, axis=2)
+    true_predictions = [[label_list[p] for (p, l) in zip(prediction, label) if l != -100]
+                        for prediction, label in zip(predictions, labels)]
+    true_labels = [[label_list[l] for (p, l) in zip(prediction, label) if l != -100]
+                   for prediction, label in zip(predictions, labels)]
+    results = metric.compute(predictions=true_predictions, references=true_labels, scheme="IOB2")
+    return {
+        "precision": results["overall_precision"],
+        "recall": results["overall_recall"],
+        "f1": results["overall_f1"],
+        "accuracy": results["overall_accuracy"],
+    }
+
+# 6. Настройка и запуск Trainer
+args = TrainingArguments(
+    output_dir=SAVE_PATH,
+    eval_strategy="epoch",
+    learning_rate=LEARNING_RATE,
+    per_device_train_batch_size=BATCH_SIZE,
+    per_device_eval_batch_size=BATCH_SIZE,
+    num_train_epochs=EPOCHS,
+    weight_decay=0.01,
+    save_strategy="no",
+    report_to=[]
+)
+
+trainer = Trainer(
+    model,
+    args,
+    train_dataset=tokenized_datasets["train"],
+    eval_dataset=tokenized_datasets["validation"],
+    data_collator=DataCollatorForTokenClassification(tokenizer),
+    tokenizer=tokenizer,
+    compute_metrics=compute_metrics
+)
+
+trainer.train()
+
+# 7. Финальное сохранение
+model.save_pretrained(SAVE_PATH)
+tokenizer.save_pretrained(SAVE_PATH)