updated token error

app.py

@@ -39,14 +39,15 @@ def get_model_infos(multilingual="multilingual"):
     return df
 
 class MLMDataset(Dataset):
-    def __init__(self,sentence,tokenizer,num_samples,MLM_MASK_TOKEN,MLM_UNK_TOKEN):
+    def __init__(self,sentence,tokenizer,MLM_MASK_TOKEN,MLM_UNK_TOKEN):
         self.sentence = sentence
         self.tokenizer = tokenizer
-        self.num_samples = len(self.sentence) - 2
 
         self.tensor_input = self.tokenizer(sentence, return_tensors='pt')['input_ids']
-        self.batch_input = self.tensor_input.repeat(self.num_samples, 1)
 
+        self.num_samples = self.tensor_input.size()[-1] - 2
+
+        self.batch_input = self.tensor_input.repeat(self.num_samples, 1)
         self.random_ids = np.random.choice([i for i in range(1,self.tensor_input.size(1)-1)],self.num_samples,replace=False) # ensuring that the masking is not done on the BOS and EOS tokens since they are not connected to the sentence itself.
         self.random_ids = torch.Tensor(self.random_ids).long().unsqueeze(0).T
 
@@ -77,8 +78,8 @@ class MLMDataset(Dataset):
         return self.masked_input[idx], self.mask[idx],self.labels[idx], self.unk_mask[idx]
 
 
-def get_sense_score_batched(sentence,tokenizer,model,MLM_MASK_TOKEN,MLM_UNK_TOKEN,num_samples,BATCH_SIZE):
-    mlm_dataset = MLMDataset(sentence,tokenizer,num_samples,MLM_MASK_TOKEN,MLM_UNK_TOKEN)
+def get_sense_score_batched(sentence,tokenizer,model,MLM_MASK_TOKEN,MLM_UNK_TOKEN,BATCH_SIZE):
+    mlm_dataset = MLMDataset(sentence,tokenizer,MLM_MASK_TOKEN,MLM_UNK_TOKEN)
     dataloader = DataLoader(mlm_dataset,batch_size=BATCH_SIZE)
 
     score =1 
@@ -119,7 +120,7 @@ def get_sense_score(sentence,tokenizer,model,MLM_MASK_TOKEN,MLM_UNK_TOKEN,num_sa
     
     tensor_input = tokenizer(sentence, return_tensors='pt')['input_ids']
     batch_input = tensor_input.repeat(num_samples, 1)
-
+    
     random_ids = np.random.choice([i for i in range(1,tensor_input.size(1)-1)],num_samples,replace=False) # ensuring that the masking is not done on the BOS and EOS tokens since they are not connected to the sentence itself.
     random_ids = torch.Tensor(random_ids).long().unsqueeze(0).T
 
@@ -190,7 +191,7 @@ if run:
     for index, model_id in enumerate(selected_models):
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         model = AutoModelWithLMHead.from_pretrained(model_id)
-        if model_id == 'castorini/afriberta_base':
+        if model_id.startswith("castorini"):
             tokenizer.model_max_length = 512 
         MLM_MASK_TOKEN = tokenizer.mask_token_id #[(103, '[MASK]')]
         MLM_UNK_TOKEN = tokenizer.unk_token_id