use 100 examples of test set

evaluate_data.py

@@ -1,5 +1,5 @@
 from evaluate_model import compute_metrics
-from datasets import load_from_disk
+from datasets import load_from_disk, Dataset
 from transformers import AutoTokenizer
 import os
 import pickle
@@ -22,6 +22,7 @@ tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 data_path = './data/merged_dataset/'
 
 test = load_from_disk(data_path)['test']
+test = Dataset.from_dict(test[:105])
 
 feature_path = './data/ner_feature.pickle'
 
@@ -53,15 +54,16 @@ def collate_fn(data):
     
     return input_ids, token_type_ids, attention_mask, labels
 
-loader = torch.utils.data.DataLoader(tokenized_test, batch_size=32, collate_fn=collate_fn)
+loader = torch.utils.data.DataLoader(tokenized_test, batch_size=16, collate_fn=collate_fn)
 device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
-print(device)
 
 ner_model = ner_model.eval()
 
 
 
 def get_metrics_trf():
+    print(device)
+    
     y_true, logits = [], []
 
     for input_ids, token_type_ids, attention_mask, labels in tqdm(loader):