| """
|
| Utility functions for tokenizer-related operations.
|
| """
|
| import torch
|
| import logging
|
| from typing import Dict, List, Any, Union, Optional
|
| from transformers import AutoTokenizer
|
|
|
| logger = logging.getLogger(__name__)
|
|
|
| def get_special_tokens_mask(tokenizer, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
|
| """
|
| Retrieve special tokens mask.
|
|
|
| Args:
|
| tokenizer: Tokenizer to use
|
| token_ids_0: First token IDs
|
| token_ids_1: Second token IDs (for pairs)
|
| already_has_special_tokens: Whether token_ids already contain special tokens
|
|
|
| Returns:
|
| List of 1s and 0s, where 1 indicates a special token
|
| """
|
| if already_has_special_tokens:
|
| return tokenizer.get_special_tokens_mask(
|
| token_ids_0,
|
| token_ids_1=token_ids_1,
|
| already_has_special_tokens=True
|
| )
|
|
|
| if token_ids_1 is None:
|
| return tokenizer.get_special_tokens_mask(
|
| token_ids_0,
|
| token_ids_1=None,
|
| already_has_special_tokens=False
|
| )
|
|
|
| return tokenizer.get_special_tokens_mask(
|
| token_ids_0,
|
| token_ids_1=token_ids_1,
|
| already_has_special_tokens=False
|
| )
|
|
|
| def add_tokens_to_tokenizer(tokenizer, new_tokens):
|
| """
|
| Add new tokens to tokenizer vocabulary.
|
|
|
| Args:
|
| tokenizer: Tokenizer to modify
|
| new_tokens: List of new tokens to add
|
|
|
| Returns:
|
| Number of tokens added
|
| """
|
| return tokenizer.add_tokens(new_tokens)
|
|
|
| def format_batch_for_model(
|
| batch: Dict[str, torch.Tensor],
|
| device: torch.device = None
|
| ) -> Dict[str, torch.Tensor]:
|
| """
|
| Format a batch for model input, moving tensors to specified device.
|
|
|
| Args:
|
| batch: Dictionary of tensors
|
| device: Device to move tensors to
|
|
|
| Returns:
|
| Formatted batch dictionary
|
| """
|
| if device is None:
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
| formatted_batch = {}
|
| for k, v in batch.items():
|
| if isinstance(v, torch.Tensor):
|
| formatted_batch[k] = v.to(device)
|
| else:
|
| formatted_batch[k] = v
|
| return formatted_batch
|
|
|
| def batch_encode_plus(
|
| tokenizer,
|
| texts: List[str],
|
| batch_size: int = 32,
|
| max_length: int = 512,
|
| return_tensors: str = "pt",
|
| **kwargs
|
| ) -> List[Dict[str, torch.Tensor]]:
|
| """
|
| Encode a large batch of texts in smaller chunks.
|
|
|
| Args:
|
| tokenizer: Tokenizer to use
|
| texts: List of texts to encode
|
| batch_size: Size of each processing batch
|
| max_length: Maximum sequence length
|
| return_tensors: Return format ('pt' for PyTorch)
|
| **kwargs: Additional encoding parameters
|
|
|
| Returns:
|
| List of encoded batches
|
| """
|
| batches = []
|
|
|
| for i in range(0, len(texts), batch_size):
|
| batch_texts = texts[i:i + batch_size]
|
| encoded = tokenizer(
|
| batch_texts,
|
| max_length=max_length,
|
| padding="max_length",
|
| truncation=True,
|
| return_tensors=return_tensors,
|
| **kwargs
|
| )
|
| batches.append(encoded)
|
|
|
| return batches
|
|
|
| def get_tokenizer_info(tokenizer) -> Dict[str, Any]:
|
| """
|
| Get information about a tokenizer.
|
|
|
| Args:
|
| tokenizer: Tokenizer to inspect
|
|
|
| Returns:
|
| Dictionary with tokenizer information
|
| """
|
| info = {
|
| "vocab_size": len(tokenizer),
|
| "model_name": getattr(tokenizer, "name_or_path", None),
|
| "special_tokens": {}
|
| }
|
|
|
|
|
| special_tokens = [
|
| "pad_token", "unk_token", "sep_token",
|
| "cls_token", "mask_token", "bos_token", "eos_token"
|
| ]
|
|
|
| for token_name in special_tokens:
|
| token_value = getattr(tokenizer, f"{token_name}", None)
|
| if token_value is not None:
|
| info["special_tokens"][token_name] = token_value
|
|
|
| return info
|
|
|
