Added model saving

app/train.py

@@ -1,10 +1,13 @@
 #!/usr/bin/env python3
 
 import numpy as np
+import tensorflow as tf
 from datasets import load_dataset
 from tensorflow.keras.optimizers import Adam
 from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
 
+checkpoint_path = "out/cp.ckpt"
+
 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)}
@@ -46,5 +49,10 @@ tf_dataset = model.prepare_tf_dataset(
     encoded["train"], batch_size=16, shuffle=True, tokenizer=tokenizer
 )
 
+cp_callback = tf.keras.callbacks.ModelCheckpoint(
+    filepath=checkpoint_path, save_weights_only=True, verbose=1
+)
+
 model.compile(optimizer=Adam(3e-5), loss="categorical_crossentropy")
-model.fit(tf_dataset)
+model.fit(tf_dataset, callbacks=[cp_callback])
+model.save('out/model')