#!/bin/bash

PLM_DIR=/home/shared/bert-base-uncased-swahili
MODEL_DIR=/home/shared/bert-swahili-news-classification
TASK_NAME=swahili_news
echo "PLM_DIR: ${PLM_DIR}"
echo "MODEL_DIR: ${MODEL_DIR}"
echo "TASK_NAME: ${TASK_NAME}"

python run_flax_glue.py \
  --model_name_or_path ${PLM_DIR} \
  --task_name ${TASK_NAME} \
  --max_length 512 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --per_device_train_batch_size 4 \
  --output_dir ${MODEL_DIR}