minBERT / constants.py
GlowCheese's picture
Hotfix
27888bb
raw
history blame contribute delete
985 Bytes
import os
DATA_DIR = 'minbert-data'
MODEL_DIR = 'minbert-model'
# Pretrained weights
SUP_BERT = os.path.join(MODEL_DIR, 'sup-cse-bert.pth')
UNSUP_BERT = os.path.join(MODEL_DIR, 'unsup-cse-bert.pth')
# CFIMDB dataset
IDS_CFIMDB_DEV = os.path.join(DATA_DIR, 'ids-cfimdb-dev.csv')
IDS_CFIMDB_TEST = os.path.join(DATA_DIR, 'ids-cfimdb-test-student.csv')
IDS_CFIMDB_TRAIN = os.path.join(DATA_DIR, 'ids-cfimdb-train.csv')
# SST dataset
IDS_SST_DEV = os.path.join(DATA_DIR, 'ids-sst-dev.csv')
IDS_SST_TEST = os.path.join(DATA_DIR, 'ids-sst-test-student.csv')
IDS_SST_TRAIN = os.path.join(DATA_DIR, 'ids-sst-train.csv')
# SimCSE train/dev dataset
NLI_TRAIN = os.path.join(DATA_DIR, 'nli-train.parquet')
AMAZON_POLARITY = os.path.join(DATA_DIR, 'amazon-polarity.parquet')
STSB_DEV = os.path.join(DATA_DIR, 'stsb-dev.parquet')
# Training-specific constants
SEED=11711
NUM_CPU_CORES=4
EPOCHS=10
USE_GPU=True
BATCH_SIZE_CSE=8
BATCH_SIZE_SST=64
BATCH_SIZE_CFIMDB=8
HIDDEN_DROPOUT_PROB=0.3