Training tweaks

app/train.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
-import numpy as np
+import multiprocessing
+
 import tensorflow as tf
 from datasets import load_dataset
 from tensorflow.keras.optimizers import Adam
@@ -18,7 +19,7 @@ label2id = {label: id for id, label in enumerate(labels)}
 id2label = {id: label for id, label in enumerate(labels)}
 
 dataset = load_dataset("csv", data_files="train.csv")
-tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
 
 def process_data(row):
@@ -27,19 +28,19 @@ def process_data(row):
 
     encoding = tokenizer(text, padding="max_length", truncation=True)
 
-    labels_matrix = np.zeros((len(text), len(labels)))
+    label_arr = [0] * len(labels)
 
     # fill numpy array
-    for id, label in enumerate(labels):
-        labels_matrix[:, id] = labels_batch[label]
+    for id, label in enumerate(labels_batch):
+        label_arr[id] = labels_batch[label]
 
-    encoding["labels"] = labels_matrix.tolist()
+    encoding["labels"] = label_arr
 
     return encoding
 
 
 model = TFAutoModelForSequenceClassification.from_pretrained(
-    "bert-base-cased",
+    "bert-base-uncased",
     problem_type="multi_label_classification",
     num_labels=len(labels),
     label2id=label2id,
@@ -47,7 +48,9 @@ model = TFAutoModelForSequenceClassification.from_pretrained(
 )
 
 encoded = dataset.map(
-    process_data, batched=True, remove_columns=[*labels, "id", "comment_text"]
+    process_data,
+    remove_columns=["id", "comment_text"],
+    num_proc=int(multiprocessing.cpu_count()),
 )
 
 tf_dataset = model.prepare_tf_dataset(