Update modeling_bertchunker.py

modeling_bertchunker.py

@@ -48,7 +48,6 @@ class BertChunker(PreTrainedModel):
         CLS=input_ids[:,0].unsqueeze(0)
         SEP=input_ids[:,-1].unsqueeze(0)
         input_ids=input_ids[:,1:-1]
-        # model= model.to(device)
         self.eval()
         split_str_poses=[]
 
@@ -57,26 +56,28 @@ class BertChunker(PreTrainedModel):
 
         while windows_end <= input_ids.shape[1]:
             windows_end= windows_start + MAX_TOKENS-2
-            
+
             ids=torch.cat((CLS, input_ids[:,windows_start:windows_end],SEP),1)
-        
+
             ids=ids.to(self.device)
-            
+
             output=self(input_ids=ids,attention_mask=attention_mask[:,:len(ids)])
             logits = output['logits'][:, 1:-1,:]
             is_left_greater = ((logits[:,:, 0] + threshold) < logits[:,:, 1])
-            greater_rows_indices = torch.where(is_left_greater)
+            greater_rows_indices = torch.where(is_left_greater)[1].tolist()
 
             # null or not
-            if greater_rows_indices[1].numel():
-                split_token_idx = greater_rows_indices[1] + windows_start + 1
+            if len(greater_rows_indices)>0 and (not (greater_rows_indices[0] == 0 and len(greater_rows_indices)==1)):
 
-                split_str_pos=[tokens.token_to_chars(sp).start for sp in split_token_idx.tolist()]
+                
+                split_str_pos=[tokens.token_to_chars(sp + windows_start + 1).start for sp in greater_rows_indices]
 
                 split_str_poses += split_str_pos
 
-                windows_start = greater_rows_indices[1][-1] + windows_start
+                windows_start = greater_rows_indices[-1] + windows_start
+
             else:
+
                 windows_start = windows_end
 
         substrings = [text[i:j] for i, j in zip([0] + split_str_poses, split_str_poses+[len(text)])]