Successfully training

app/.gitignore

@@ -0,0 +1,2 @@
+out/
+train.csv

app/train.py

@@ -1,36 +1,34 @@
 #!/usr/bin/env python3
 
-from collections import defaultdict
-
-import pandas as pd
-import tensorflow as tf
-from transformers import (
-    AutoTokenizer,
-    TFAutoModelForSequenceClassification,
-    TFTrainer,
-    TFTrainingArguments,
-)
+import numpy as np
+from datasets import load_dataset
+from tensorflow.keras.optimizers import Adam
+from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
 
 labels = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
 label2id = {label: id for id, label in enumerate(labels)}
 id2label = {id: label for id, label in enumerate(labels)}
 
-data = pd.read_csv("./train.csv")
-
-batch_encodings = defaultdict(list)
-batch_labels = list()
+dataset = load_dataset("csv", data_files="train.csv")
 tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
 
-for _, row in data.sample(n=3).iterrows():
+
+def process_data(row):
     text = row["comment_text"]
-    text_labels = {k: row[k] for k in row.keys() if k in labels}
+    labels_batch = {k: row[k] for k in row.keys() if k in labels}
 
     encoding = tokenizer(text, padding="max_length", truncation=True)
 
-    batch_labels.append([text_labels[id2label[id]] for id in range(len(labels))])
+    labels_matrix = np.zeros((len(text), len(labels)))
+
+    # fill numpy array
+    for id, label in enumerate(labels):
+        labels_matrix[:, id] = labels_batch[label]
+
+    encoding["labels"] = labels_matrix.tolist()
+
+    return encoding
 
-    for key in encoding.keys():
-        batch_encodings[key].append(encoding[key])
 
 model = TFAutoModelForSequenceClassification.from_pretrained(
     "bert-base-cased",
@@ -40,15 +38,13 @@ model = TFAutoModelForSequenceClassification.from_pretrained(
     id2label=id2label,
 )
 
-training_args = TFTrainingArguments(output_dir="test_trainer")
-train_dataset = tf.data.Dataset.from_tensor_slices(
-    (dict(batch_encodings), batch_labels)
+encoded = dataset.map(
+    process_data, batched=True, remove_columns=[*labels, "id", "comment_text"]
 )
 
-trainer = TFTrainer(
-    model=model,
-    args=training_args,
-    train_dataset=train_dataset,
+tf_dataset = model.prepare_tf_dataset(
+    encoded["train"], batch_size=16, shuffle=True, tokenizer=tokenizer
 )
 
-trainer.train()
+model.compile(optimizer=Adam(3e-5), loss="categorical_crossentropy")
+model.fit(tf_dataset)