Upload hybridna_tokenizer.py with huggingface_hub

hybridna_tokenizer.py

@@ -160,6 +160,8 @@ class HybriDNATokenizer(PreTrainedTokenizer):
         truncation: bool = True,
         max_length: Optional[int] = None,
         add_special_tokens: bool = True,
+        return_tensors: Optional[str] = None,
+        **kwargs,
     ):
         # ---------- detect batch vs single ----------
         is_batch = not isinstance(text, str)
@@ -220,7 +222,26 @@ class HybriDNATokenizer(PreTrainedTokenizer):
             batch_input_ids = batch_input_ids[0]
             batch_attention = batch_attention[0]
 
-        return {
+        result = {
             "input_ids": batch_input_ids,
             "attention_mask": batch_attention,
         }
+
+        # ---------- convert to tensors if requested ----------
+        if return_tensors == "pt":
+            import torch
+            if is_batch:
+                result["input_ids"] = torch.tensor(np.stack(result["input_ids"]), dtype=torch.long)
+                result["attention_mask"] = torch.tensor(np.stack(result["attention_mask"]), dtype=torch.long)
+            else:
+                result["input_ids"] = torch.tensor(result["input_ids"], dtype=torch.long).unsqueeze(0)
+                result["attention_mask"] = torch.tensor(result["attention_mask"], dtype=torch.long).unsqueeze(0)
+        elif return_tensors == "np":
+            if is_batch:
+                result["input_ids"] = np.stack(result["input_ids"])
+                result["attention_mask"] = np.stack(result["attention_mask"])
+            else:
+                result["input_ids"] = np.expand_dims(result["input_ids"], 0)
+                result["attention_mask"] = np.expand_dims(result["attention_mask"], 0)
+
+        return result