feat(model): add stats

dvc.lock

@@ -10,8 +10,8 @@ stages:
       md5: 465d64e91ac7d713e57442c932a665ea
       size: 2889
     - path: shad_mlops_transformers/trainer.py
-      md5: 39d02f490890ef306b1165fcc29b8895
-      size: 7495
+      md5: 6f76b90e32a94bea31bc9365494aceb9
+      size: 7503
     params:
       shad_mlops_transformers/config.py:
         Config.batch_size: 64
@@ -20,11 +20,11 @@ stages:
         Config.test_size: 0.2
     outs:
     - path: data/models/model.torch
-      md5: cad526c49e9a9ef312e820ba4cc93442
-      size: 438842393
+      md5: 555452b7e7eedffff70cae47c8a1e2f5
+      size: 438914457
     - path: data/plots/train.png
-      md5: a23fefe43ee875860f875123bd402522
-      size: 40498
+      md5: 6ea1333efcffb58aa7f3b9161a3cb488
+      size: 35617
     - path: data/raw/mapper.json
-      md5: e4c22adf587676d22ed58c812f41f409
-      size: 855
+      md5: 47758bce6ee09369e716fd8658ce3e01
+      size: 930

shad_mlops_transformers/trainer.py

@@ -20,7 +20,7 @@ class ArxivDataset(Dataset):
         self.x = []
         self.y = []
         # self.data = []
-        whitelist_labels = ["math", "cs"]
+        whitelist_labels = ["math", "cs", "stat"]
         i = 0
         if class_mapper is None:
             self.class_mapper = {}