Upload minitransformer.py with huggingface_hub

minitransformer.py

@@ -7,7 +7,7 @@ import os
 from transformers import AutoTokenizer, logging
 import pandas as pd
 from tqdm import tqdm
-
+from safetensors.torch import save_file
 
 logging.set_verbosity_error()
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -15,14 +15,14 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 # ----------------- CONFIG -----------------
 SAVE_EVERY = 5
 MODEL_NAME = "mini_transformer_v3"
-N_DATA_WORKERS = 6
+N_DATA_WORKERS = 8
 PIN_MEMORY = True if N_DATA_WORKERS > 0 and torch.cuda.is_available() else False
-BATCH_SIZE = 128
+BATCH_SIZE = 512
 EVAL_EVERY = 5
 LEARNING_RATE = 3e-4
 NUM_EPOCHS = 50
 USE_AMP = True
-STRIDE = 32
+STRIDE = 64
 CHECKPOINT_DIR = f"MODELS/checkpoints/{MODEL_NAME}"
 os.makedirs(CHECKPOINT_DIR, exist_ok=True)
 DATASET = "DATA/generated_dataset_very_big.csv"
@@ -263,6 +263,11 @@ for epoch in range(start_epoch, NUM_EPOCHS):
             },
             os.path.join(CHECKPOINT_DIR, f"checkpoint_{MODEL_NAME}_epoch_{epoch+1}.pt"),
         )
+        save_file(
+            model.state_dict(),
+            os.path.join(CHECKPOINT_DIR, f"model_{epoch+1}.safetensors"),
+        )
+
 
 # check GPU utilization metrics here:
 # nvidia-smi dmon -s u