Override _encode_sequences and _predict_logits to use HuggingFace tokenizer interface

adapter.py

@@ -191,12 +191,13 @@ class AbLang2PairedHuggingFaceAdapter(AbEncoding, AbRestore, AbAlignment, AbScor
         self.AbLang.train()
 
     def _encode_sequences(self, seqs):
-        # Use HuggingFace-style padding and return PyTorch tensors
+        # Override to use HuggingFace tokenizer interface
         tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
         tokens = extract_input_ids(tokens, self.used_device)
         return self.AbRep(tokens).last_hidden_states.detach()
 
     def _predict_logits(self, seqs):
+        # Override to use HuggingFace tokenizer interface
         tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
         tokens = extract_input_ids(tokens, self.used_device)
         output = self.AbLang(tokens)
@@ -204,6 +205,26 @@ class AbLang2PairedHuggingFaceAdapter(AbEncoding, AbRestore, AbAlignment, AbScor
             return output.last_hidden_state.detach()
         return output.detach()
 
+    def _predict_logits_with_step_masking(self, seqs):
+        # Override the stepwise masking method to use HuggingFace tokenizer
+        tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
+        tokens = extract_input_ids(tokens, self.used_device)
+        
+        logits = []
+        for single_seq_tokens in tokens:
+            tkn_len = len(single_seq_tokens)
+            masked_tokens = single_seq_tokens.repeat(tkn_len, 1)
+            for num in range(tkn_len):
+                masked_tokens[num, num] = self.tokenizer.mask_token_id
+            
+            with torch.no_grad():
+                logits_tmp = self.AbLang(masked_tokens)
+                       
+            logits_tmp = torch.stack([logits_tmp[num, num] for num in range(tkn_len)])
+            logits.append(logits_tmp)
+    
+        return torch.stack(logits, dim=0)
+
     def _preprocess_labels(self, labels):
         labels = extract_input_ids(labels, self.used_device)
         return labels