Create train.py

train.py

@@ -0,0 +1,73 @@
+import pandas as pd
+import torch
+from transformers import DebertaV2ForSequenceClassification, DebertaV2Tokenizer, DataCollatorWithPadding, Trainer, TrainingArguments
+from tqdm import tqdm
+from datasets import Dataset, load_dataset
+import numpy as np
+import wandb
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
+
+output_dir = './german_politic_DeBERTa-v2-base'
+model_name = "ikim-uk-essen/geberta-base"
+max_length = 512
+id2label = {0: 'other', 1: 'politic'}
+label2id = {'other': 0, 'politic': 1}
+
+wandb.init(project="german_politic_yes_no_classifier", entity="xxx", name="german_politic_DeBERTa")
+
+model = DebertaV2ForSequenceClassification.from_pretrained(model_name, num_labels = 2, id2label=id2label, label2id=label2id, output_attentions = False, output_hidden_states = False)
+tokenizer = DebertaV2Tokenizer.from_pretrained(model_name, do_lower_case=False, max_length = max_length, TOKENIZERS_PARALLELISM=True)
+
+dataset = load_dataset("SinclairSchneider/trainset_political_text_yes_no_german")
+dataset = dataset['train'].train_test_split(0.2)
+
+def preprocess(sample):
+    return tokenizer(sample["text"], truncation=True)
+
+dataset_tokenized = dataset.map(preprocess, batched = True)
+
+data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
+def compute_metrics(pred):
+    labels = pred.label_ids
+    preds = pred.predictions.argmax(-1)
+    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='weighted')
+    acc = accuracy_score(labels, preds)
+    return {
+        'accuracy': acc,
+        'f1': f1,
+        'precision': precision,
+        'recall': recall
+    }
+
+training_args = TrainingArguments(
+    output_dir = output_dir,
+    learning_rate=2e-5,
+    per_device_train_batch_size=16,
+    per_device_eval_batch_size=16,
+    num_train_epochs=4,
+    weight_decay=0.01,
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    load_best_model_at_end=True,
+    report_to="wandb",
+    fp16 = False,
+    logging_steps = 8,
+    disable_tqdm = False,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset_tokenized["train"],
+    eval_dataset=dataset_tokenized["test"],
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+)
+
+trainer.train()
+
+model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training
+model_to_save.save_pretrained(output_dir)
+tokenizer.save_pretrained(output_dir)