ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: 修改模型架构

Browse Source
This commit is contained in:
ViperEkura
2026-04-16 18:57:21 +08:00
parent cb93f8219e
commit a38334f4ce
17 changed files with 540 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pipeline/__init__.py
+5 -2
View File
@@ -1,5 +1,5 @@
import logging
from pipeline.tokenizer import BpeTokenizer
from pipeline.tokenize import AutoTokenizer, ChatTemplate, train_bpe_tokenizer
from pipeline.text import TextNormalizer
from pipeline.packing import SequencePacker
from pipeline.io import IOHandler, export_dataset, cache_jsonl
@@ -16,8 +16,11 @@ from pipeline.strategies import (
setup_logging()
__all__ = [
# Tokenizer
"AutoTokenizer",
"ChatTemplate",
"train_bpe_tokenizer",
# Core modules
"BpeTokenizer",
"TextNormalizer",
"SequencePacker",
"IOHandler",
pipeline/processors/dpo.py
+2 -2
View File
@@ -5,7 +5,7 @@ from typing import Dict, List, Any, Optional
import torch
from torch import Tensor
from pipeline.tokenizer import BpeTokenizer
from pipeline.tokenize import AutoTokenizer
from pipeline.strategies import PromptStrategy, ChatMLStrategy
from pipeline.processors.base import BaseProcessor, _encode_with_mask
from pipeline.processors.factory import ProcessorFactory
@@ -20,7 +20,7 @@ class DPOProcessor(BaseProcessor):
def __init__(
self,
tokenizer: BpeTokenizer,
tokenizer: AutoTokenizer,
strategy: Optional[PromptStrategy] = None,
):
self.tokenizer = tokenizer
pipeline/processors/factory.py
+4 -4
View File
@@ -3,7 +3,7 @@
from typing import Dict, List, Any, Optional, Type
from pipeline.processors.base import BaseProcessor
from pipeline.tokenizer import BpeTokenizer
from pipeline.tokenize import AutoTokenizer
from pipeline.strategies import PromptStrategy, StrategyFactory
@@ -45,7 +45,7 @@ class ProcessorFactory:
return decorator
@classmethod
def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> BaseProcessor:
def create(cls, processor_type: str, tokenizer: AutoTokenizer) -> BaseProcessor:
"""Create a processor by type name (uses default ChatMLStrategy for SFT/DPO).
Args:
@@ -69,7 +69,7 @@ class ProcessorFactory:
def create_with_strategy(
cls,
processor_type: str,
tokenizer: BpeTokenizer,
tokenizer: AutoTokenizer,
strategy: PromptStrategy,
) -> BaseProcessor:
"""Create a processor with a custom strategy.
@@ -99,7 +99,7 @@ class ProcessorFactory:
def create_with_strategy_name(
cls,
processor_type: str,
tokenizer: BpeTokenizer,
tokenizer: AutoTokenizer,
strategy_name: str,
**strategy_kwargs,
) -> BaseProcessor:
pipeline/processors/pretrain.py
+2 -2
View File
@@ -5,7 +5,7 @@ from typing import Dict, List, Any
import torch
from torch import Tensor
from pipeline.tokenizer import BpeTokenizer
from pipeline.tokenize import AutoTokenizer
from pipeline.processors.base import BaseProcessor
from pipeline.processors.factory import ProcessorFactory
@@ -14,7 +14,7 @@ from pipeline.processors.factory import ProcessorFactory
class PreTrainProcessor(BaseProcessor):
"""Pre-training data processor."""
def __init__(self, tokenizer: BpeTokenizer):
def __init__(self, tokenizer: AutoTokenizer):
self.tokenizer = tokenizer
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
pipeline/processors/sft.py
+2 -2
View File
@@ -5,7 +5,7 @@ from typing import Dict, List, Any, Optional
import torch
from torch import Tensor
from pipeline.tokenizer import BpeTokenizer
from pipeline.tokenize import AutoTokenizer
from pipeline.strategies import PromptStrategy, ChatMLStrategy
from pipeline.processors.base import BaseProcessor, _encode_with_mask
from pipeline.processors.factory import ProcessorFactory
@@ -20,7 +20,7 @@ class SFTProcessor(BaseProcessor):
def __init__(
self,
tokenizer: BpeTokenizer,
tokenizer: AutoTokenizer,
strategy: Optional[PromptStrategy] = None,
):
self.tokenizer = tokenizer
pipeline/strategies/alpaca.py
+2 -2
View File
@@ -2,7 +2,7 @@
from typing import List
from pipeline.tokenizer import BpeTokenizer
from pipeline.tokenize import AutoTokenizer
from pipeline.strategies.base import PromptStrategy
from pipeline.strategies.factory import StrategyFactory
@@ -13,7 +13,7 @@ class AlpacaStrategy(PromptStrategy):
def __init__(
self,
tokenizer: BpeTokenizer,
tokenizer: AutoTokenizer,
instruction_start: str = "### Instruction:\n",
response_start: str = "### Response:\n",
response_suffix: str = "\n<｜end▁of▁sentence｜>",
pipeline/strategies/base.py
+2 -2
View File
@@ -3,7 +3,7 @@
from abc import ABC, abstractmethod
from typing import List
from pipeline.tokenizer import BpeTokenizer
from pipeline.tokenize import AutoTokenizer
class PromptStrategy(ABC):
@@ -14,7 +14,7 @@ class PromptStrategy(ABC):
which assembles them with pre-encoded format tokens.
"""
def __init__(self, tokenizer: BpeTokenizer):
def __init__(self, tokenizer: AutoTokenizer):
self.tokenizer = tokenizer
def _encode_format(self, text: str) -> List[int]:
pipeline/strategies/chatml.py
+2 -2
View File
@@ -2,7 +2,7 @@
from typing import List
from pipeline.tokenizer import BpeTokenizer
from pipeline.tokenize import AutoTokenizer
from pipeline.strategies.base import PromptStrategy
from pipeline.strategies.factory import StrategyFactory
@@ -13,7 +13,7 @@ class ChatMLStrategy(PromptStrategy):
def __init__(
self,
tokenizer: BpeTokenizer,
tokenizer: AutoTokenizer,
user_start: str = "<｜im▁start｜>user\n",
user_end: str = "<｜im▁end｜>\n",
assistant_start: str = "<｜im▁start｜>assistant\n",
pipeline/strategies/factory.py
+2 -2
View File
@@ -2,7 +2,7 @@
from typing import Dict, List, Type
from pipeline.tokenizer import BpeTokenizer
from pipeline.tokenize import AutoTokenizer
from pipeline.strategies.base import PromptStrategy
@@ -44,7 +44,7 @@ class StrategyFactory:
return decorator
@classmethod
def create(cls, name: str, tokenizer: BpeTokenizer, **kwargs) -> PromptStrategy:
def create(cls, name: str, tokenizer: AutoTokenizer, **kwargs) -> PromptStrategy:
"""Create a strategy by name.
Args:
pipeline/tokenize/__init__.py
+12
View File
@@ -0,0 +1,12 @@
"""
Tokenizer module with BPE implementation and auto-loading support.
"""
from pipeline.tokenize.tokenizer import AutoTokenizer, train_bpe_tokenizer
from pipeline.tokenize.chat_template import ChatTemplate
__all__ = [
"AutoTokenizer",
"train_bpe_tokenizer",
"ChatTemplate",
]
pipeline/tokenize/chat_template.py
+120
View File
@@ -0,0 +1,120 @@
"""
Chat template module with Jinja2 rendering support.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from jinja2 import Template
# Message type for chat messages
type MessageType = Dict[str, Any]
@dataclass
class ChatTemplate:
"""A chat template with Jinja2 rendering support.
Attributes:
name: Unique identifier for the template.
template_str: Jinja2 template string.
description: Optional description.
default_variables: Optional dictionary of default variable values
that will be passed to the template if not overridden during rendering.
special_tokens: Optional dictionary mapping token names to their string values.
These tokens are automatically added to the template variables.
"""
name: str
template_str: str
description: str = ""
default_variables: Dict[str, Any] = field(default_factory=dict)
special_tokens: Dict[str, str] = field(default_factory=dict)
@classmethod
def from_string(
cls,
template_str: str,
description: str = "",
default_variables: Optional[Dict[str, Any]] = None,
special_tokens: Optional[Dict[str, str]] = None,
) -> "ChatTemplate":
"""Create a ChatTemplate instance directly from a template string."""
return cls(
name="", # empty name for ad-hoc templates
template_str=template_str,
description=description,
default_variables=default_variables or {},
special_tokens=special_tokens or {},
)
def render(
self,
messages: List[MessageType],
system_prompt: Optional[str] = None,
add_generation_prompt: bool = True,
**extra_variables: Any,
) -> str:
"""Render the template with given messages and variables.
Args:
messages: List of message dicts with 'role' and 'content'.
system_prompt: Optional system prompt string.
add_generation_prompt: Whether to add generation prompt after messages.
**extra_variables: Additional variables to pass to the template.
These override default_variables and special_tokens.
Returns:
Rendered prompt string.
"""
# Merge default variables, special tokens, and extra variables
variables = {
**self.default_variables,
**self.special_tokens,
**extra_variables,
}
variables["messages"] = messages
variables["add_generation_prompt"] = add_generation_prompt
if system_prompt is not None:
variables["system_prompt"] = system_prompt
jinja_template = Template(self.template_str)
return jinja_template.render(**variables)
# Default ChatML template
DEFAULT_CHATML_TEMPLATE = """{% for message in messages %}{{ bos_token }}{{ message['role'] }}
{{ message['content'] }}{{ eos_token }}{% endfor %}{% if add_generation_prompt %}{{ bos_token }}assistant
{% endif %}"""
# Pre-built template registry
TEMPLATE_REGISTRY: Dict[str, ChatTemplate] = {}
def register_chat_template(name: str, template: ChatTemplate) -> None:
"""Register a chat template in the global registry.
Args:
name: Name to register the template under
template: ChatTemplate instance
"""
TEMPLATE_REGISTRY[name] = template
def get_chat_template(name: str) -> ChatTemplate:
"""Get a registered chat template.
Args:
name: Template name
Returns:
ChatTemplate instance
Raises:
KeyError: If template not found
"""
if name not in TEMPLATE_REGISTRY:
raise KeyError(f"Chat template '{name}' not found. Available: {list(TEMPLATE_REGISTRY.keys())}")
return TEMPLATE_REGISTRY[name]
pipeline/tokenize/tokenizer.py
+378
View File
@@ -0,0 +1,378 @@
"""
Tokenizer module with BPE implementation and auto-loading support.
"""
from dataclasses import dataclass
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
from tokenizers import Tokenizer
from tokenizers import decoders, processors, normalizers, pre_tokenizers
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from jinja2 import Template
DEFAULT_SPECIAL_TOKENS = {
"bos_token": "<｜begin▁of▁sentence｜>",
"eos_token": "<｜end▁of▁sentence｜>",
"pad_token": "<｜▁pad▁｜>",
}
CONTROL_TOKENS = [
"<｜begin▁of▁sentence｜>",
"<｜end▁of▁sentence｜>",
"<｜▁pad▁｜>",
]
SPECIAL_TOKENS = ["<｜im▁start｜>", "<｜im▁end｜>"]
def train_bpe_tokenizer(
files: List[str],
vocab_size: int,
min_freq: int = 2,
reserved_token_size: int = 100,
max_token_length: int = 18,
) -> Tokenizer:
reserved_tokens = [
f"<｜reserve{i:02d}｜>"
for i in range(reserved_token_size - len(SPECIAL_TOKENS))
]
detail_vocab_size = vocab_size - (len(reserved_tokens) + len(SPECIAL_TOKENS))
alphabet = pre_tokenizers.ByteLevel.alphabet()
min_size = len(alphabet) + len(CONTROL_TOKENS)
assert detail_vocab_size > min_size
tokenizer = Tokenizer(BPE())
tokenizer.normalizer = normalizers.Sequence([normalizers.NFC(), normalizers.Strip()])
tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
pre_tokenizers.UnicodeScripts(),
pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True),
])
tokenizer.decoder = decoders.ByteLevel()
tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
trainer = BpeTrainer(
vocab_size=detail_vocab_size,
min_frequency=min_freq,
limit_alphabet=detail_vocab_size // 6,
max_token_length=max_token_length,
special_tokens=CONTROL_TOKENS + SPECIAL_TOKENS,
initial_alphabet=alphabet,
show_progress=True,
)
tokenizer.train(files=files, trainer=trainer)
tokenizer.add_special_tokens(CONTROL_TOKENS + SPECIAL_TOKENS + reserved_tokens)
return tokenizer
# Message type for chat messages
type MessageType = Dict[str, Any]
@dataclass
class ChatTemplate:
"""A chat template with Jinja2 rendering support.
Attributes:
name: Unique identifier for the template.
template_str: Jinja2 template string.
description: Optional description.
default_variables: Optional dictionary of default variable values
that will be passed to the template if not overridden during rendering.
special_tokens: Optional dictionary mapping token names to their string values.
These tokens are automatically added to the template variables.
"""
name: str
template_str: str
description: str = ""
default_variables: Dict[str, Any] = None
special_tokens: Dict[str, str] = None
def __post_init__(self):
if self.default_variables is None:
self.default_variables = {}
if self.special_tokens is None:
self.special_tokens = {}
@classmethod
def from_string(
cls,
template_str: str,
description: str = "",
default_variables: Optional[Dict[str, Any]] = None,
special_tokens: Optional[Dict[str, str]] = None,
) -> "ChatTemplate":
"""Create a ChatTemplate instance directly from a template string."""
return cls(
name="", # empty name for ad‑hoc templates
template_str=template_str,
description=description,
default_variables=default_variables,
special_tokens=special_tokens,
)
def render(
self,
messages: List[MessageType],
system_prompt: Optional[str] = None,
**extra_variables: Any,
) -> str:
"""Render the template with given messages and variables.
Args:
messages: List of message dicts with 'role' and 'content'.
system_prompt: Optional system prompt string.
**extra_variables: Additional variables to pass to the template.
These override default_variables and special_tokens.
Returns:
Rendered prompt string.
"""
# Merge default variables, special tokens, and extra variables
variables = {**self.default_variables, **self.special_tokens, **extra_variables}
variables["messages"] = messages
if system_prompt is not None:
variables["system_prompt"] = system_prompt
jinja_template = Template(self.template_str)
return jinja_template.render(**variables)
class AutoTokenizer:
"""Base tokenizer class with automatic loading support"""
TOKENIZER_CLASSES = {} # Registry for auto-loading
def __init__(
self,
path: Optional[Union[str, Path]] = None,
special_token_map: Optional[Dict[str, str]] = None,
chat_template: Optional[str] = None,
):
self._tokenizer: Tokenizer = None
self._chat_template: Optional[ChatTemplate] = None
self._special_token_map: Optional[Dict] = special_token_map or {}
if chat_template:
self.set_chat_template(chat_template)
if path:
self.load(path)
def load(self, path: Union[str, Path]):
"""Load tokenizer from directory."""
path = Path(path)
tokenizer_file = path / "tokenizer.json"
config_file = path / "tokenizer_config.json"
self._tokenizer = Tokenizer.from_file(str(tokenizer_file))
if config_file.exists():
with open(config_file, "r", encoding="utf-8") as f:
config = json.load(f)
if "special_tokens" in config:
self._special_token_map.update(config["special_tokens"])
# Load chat template from config
if "chat_template" in config:
self.set_chat_template(config["chat_template"])
@classmethod
def from_pretrained(cls, path: Union[str, Path], **kwargs) -> "AutoTokenizer":
"""Load tokenizer from pretrained directory."""
instance = cls(path)
return instance
def save_pretrained(self, save_path: str):
"""
Save tokenizer to pretrained directory.
Args:
save_path: Path to save the tokenizer
"""
save_path = Path(save_path)
save_path.mkdir(parents=True, exist_ok=True)
# Save tokenizer
self._tokenizer.save(str(save_path / "tokenizer.json"))
# Save tokenizer config
config = {}
if self._special_token_map is not None:
config["special_tokens"] = self._special_token_map
if self._chat_template is not None:
config["chat_template"] = self._chat_template.template_str
with open(save_path / "tokenizer_config.json", "w", encoding="utf-8") as f:
json.dump(config, f, ensure_ascii=False, indent=2)
@classmethod
def register_tokenizer(cls, name: str, tokenizer_class: type):
"""
Register a new tokenizer class.
Args:
name: Name to register the tokenizer class under
tokenizer_class: The tokenizer class to register
"""
cls.TOKENIZER_CLASSES[name] = tokenizer_class
def encode(
self,
tokens: Union[str, List[str]],
out_ids: bool = True,
is_pretokenized: bool = False,
add_special_tokens: bool = True,
) -> List:
"""Encode text to tokens or token IDs."""
if self._tokenizer is None:
raise RuntimeError(
"Tokenizer not initialized. Load or create a tokenizer first."
)
if isinstance(tokens, str):
encoded = self._tokenizer.encode(
tokens,
is_pretokenized=is_pretokenized,
add_special_tokens=add_special_tokens,
)
return encoded.ids if out_ids else encoded.tokens
else:
encoded_list = self._tokenizer.encode_batch(
tokens,
is_pretokenized=is_pretokenized,
add_special_tokens=add_special_tokens,
)
return [
encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
]
def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
"""Decode token IDs to text."""
if self._tokenizer is None:
raise RuntimeError(
"Tokenizer not initialized. Load or create a tokenizer first."
)
return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
def __len__(self) -> int:
if self._tokenizer is None:
return 0
return self._tokenizer.get_vocab_size()
def __getattr__(self, key: str):
"""
Dynamically intercept special token attribute access.
Supports three forms:
- tokenizer.bos_token → returns string
- tokenizer.bos_token_id → returns corresponding integer ID
- tokenizer.stop_ids → returns list of corresponding integer IDs for all special tokens
"""
# Handle stop_ids - return IDs for all special tokens
if key == "stop_ids":
stop_ids = []
if self._tokenizer is None:
return stop_ids
for val in self._special_token_map.values():
token_id = self._tokenizer.token_to_id(val)
if token_id is not None:
stop_ids.append(token_id)
return stop_ids
# Handle _id suffix (e.g., bos_token_id -> bos_token)
if key.endswith("_id"):
base_attr = key[:-3] # Remove "_id"
token_str = self._special_token_map.get(base_attr)
if token_str is None:
return None
if self._tokenizer is None:
raise RuntimeError("Tokenizer not loaded, cannot convert token to id.")
return self._tokenizer.token_to_id(token_str)
# Handle regular string attributes
if key in self._special_token_map:
return self._special_token_map.get(key)
# Other attributes trigger default AttributeError
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{key}'")
@property
def vocab_size(self) -> int:
return len(self)
def set_chat_template(self, template: Union[str, ChatTemplate]):
"""
Set the chat template for the tokenizer.
Args:
template: Either a template name (str) registered in the global registry,
or a ChatTemplate instance, or a Jinja2 template string.
Raises:
KeyError: If template name is not registered.
"""
if isinstance(template, str):
self._chat_template = ChatTemplate.from_string(template)
elif isinstance(template, ChatTemplate):
self._chat_template = template
else:
raise ValueError("Invalid template type, must be str or ChatTemplate.")
def apply_chat_template(
self,
messages: List[Dict[str, str]],
system_prompt: Optional[str] = None,
tokenize: bool = True,
add_generation_prompt: bool = True,
**kwargs,
) -> Union[str, List[int]]:
"""
Apply the chat template to messages and optionally tokenize the result.
Args:
messages: List of message dicts with 'role' and 'content'.
system_prompt: Optional system prompt string (auto-converted to first message).
tokenize: Whether to return token IDs (True) or raw string (False).
add_generation_prompt: Whether to add the generation prompt (default: True).
**kwargs: Additional variables to pass to the template.
Returns:
Either the rendered string or list of token IDs.
Raises:
RuntimeError: If chat template is not set.
"""
if self._chat_template is None:
raise RuntimeError(
"Chat template not set. Use set_chat_template() to set a template first."
)
# Auto-convert system_prompt to first message if provided
if system_prompt:
messages = [{"role": "system", "content": system_prompt}] + list(messages)
# Render the template
rendered = self._chat_template.render(
messages=messages,
add_generation_prompt=add_generation_prompt,
**kwargs,
)
if tokenize:
return self.encode(rendered)
return rendered
pipeline/tokenizer.py
-149
View File
@@ -1,149 +0,0 @@
from tokenizers import Tokenizer, Encoding
from tokenizers import decoders, processors, normalizers, pre_tokenizers
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from typing import List, Union, Optional, Tuple, Iterator
class BpeTokenizer:
def __init__(self, path: Optional[str] = None):
self._control_tokens = [
"<｜begin▁of▁sentence｜>",
"<｜end▁of▁sentence｜>",
"<｜▁pad▁｜>",
]
self._special_tokens = ["<｜im▁start｜>", "<｜im▁end｜>"]
model = BPE()
self._tokenizer = Tokenizer(model)
self._tokenizer.normalizer = normalizers.Sequence(
[normalizers.NFC(), normalizers.Strip()]
)
self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
[
pre_tokenizers.UnicodeScripts(),
pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True),
]
)
self._tokenizer.decoder = decoders.ByteLevel()
self._tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
if path is not None:
self._tokenizer = Tokenizer.from_file(path)
def _prepare_trainer(
self,
vocab_size: int,
min_freq: int,
reserved_token_size: int,
max_token_length: int = 18,
) -> Tuple[BpeTrainer, int, List[str]]:
assert reserved_token_size > len(self._special_tokens)
reserved_tokens = [
f"<｜reserve{i:02d}｜>"
for i in range(reserved_token_size - len(self._special_tokens))
]
detail_vocab_size = vocab_size - (
len(reserved_tokens) + len(self._special_tokens)
)
alphabet = pre_tokenizers.ByteLevel.alphabet()
min_size = len(alphabet) + len(self._control_tokens)
assert detail_vocab_size > min_size
trainer = BpeTrainer(
vocab_size=detail_vocab_size,
min_frequency=min_freq,
limit_alphabet=detail_vocab_size // 6,
max_token_length=max_token_length,
special_tokens=self._control_tokens + self._special_tokens,
initial_alphabet=alphabet,
show_progress=True,
)
return trainer, detail_vocab_size, reserved_tokens
def train(
self,
files: List[str],
vocab_size: int,
min_freq: int,
reserved_token_size: int = 100,
) -> None:
trainer, _, reserved_tokens = self._prepare_trainer(
vocab_size=vocab_size,
min_freq=min_freq,
reserved_token_size=reserved_token_size,
)
self._tokenizer.train(files=files, trainer=trainer)
self._tokenizer.add_special_tokens(
self._control_tokens + self._special_tokens + reserved_tokens
)
def train_from_iterator(
self,
iterator: Iterator[str],
vocab_size: int,
min_freq: int,
reserved_token_size: int = 100,
) -> None:
trainer, _, reserved_tokens = self._prepare_trainer(
vocab_size=vocab_size,
min_freq=min_freq,
reserved_token_size=reserved_token_size,
)
self._tokenizer.train_from_iterator(iterator=iterator, trainer=trainer)
self._tokenizer.add_special_tokens(
self._control_tokens + self._special_tokens + reserved_tokens
)
def save(self, path: str) -> None:
self._tokenizer.save(path)
def load(self, path: str) -> None:
self._tokenizer = Tokenizer.from_file(path)
def encode(
self,
tokens: Union[str, List[str]],
out_ids: bool = True,
add_special_tokens: bool = False,
) -> Union[List[int], List[str], List[List[int]], List[List[str]]]:
if isinstance(tokens, str):
encoded: Encoding = self._tokenizer.encode(
tokens, add_special_tokens=add_special_tokens
)
return encoded.ids if out_ids else encoded.tokens
elif isinstance(tokens, list):
encoded_list: List[Encoding] = self._tokenizer.encode_batch(
tokens, add_special_tokens=add_special_tokens
)
return [
encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
]
def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
def __len__(self) -> int:
return self._tokenizer.get_vocab_size()
@property
def stop_ids(self) -> List[int]:
stop_token = self._control_tokens + self._special_tokens
stop_ids = [self._tokenizer.token_to_id(token) for token in stop_token]
return stop_ids
@property
def bos_id(self) -> int:
return self._tokenizer.token_to_id("<｜begin▁of▁sentence｜>")
@property
def eos_id(self) -> int:
return self._tokenizer.token_to_id("<｜end▁of▁sentence｜>")
@property
def pad_id(self) -> int:
return self._tokenizer.token_to_id("<｜▁pad▁｜>")