Create example_predict.yml

dhvalden · Mar 30, 2022 · d983959 · d983959
1 parent 1be1559
commit d983959
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/example_predict.yml b/example_predict.yml
@@ -0,0 +1,13 @@
+HF:
+  model_arch: distilbert
+  model_name: distilbert-base-uncased
+CONFIG:
+  data_path: data/data_new.csv
+  text_col: 0
+  out_path: data/preds.csv
+  batch_size: 32
+  state_dict_path: models/model1_distilBERT_epoch_5.model
+  labels:
+        label1: 0
+        label2: 1
+        label3: 2