File size: 5,083 Bytes
a006332 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """BPE tokenizer training and chat formatting for Delta Ultra Mini."""
from __future__ import annotations
import logging
import os
from pathlib import Path
from typing import Any
from tokenizers import Tokenizer
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.models import BPE
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.processors import TemplateProcessing
from tokenizers.trainers import BpeTrainer
logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)
SPECIAL_TOKENS: list[str] = ["[PAD]", "[UNK]", "[BOS]", "[EOS]", "[SYS]", "[USR]", "[ASS]", "[SEP]"]
DEFAULT_SYSTEM_PROMPT = (
"Você é Delta, assistente criada pela Flame Corporation. "
"Responda de forma clara, útil e amigável."
)
def train_tokenizer(corpus_files: list[str] | list[Path], output_path: str | Path) -> None:
"""Train a BPE tokenizer from raw text files.
Args:
corpus_files: Paths to corpus files.
output_path: Destination tokenizer JSON path.
"""
tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
tokenizer.decoder = ByteLevelDecoder()
trainer = BpeTrainer(vocab_size=32000, special_tokens=SPECIAL_TOKENS, show_progress=True)
tokenizer.train([str(path) for path in corpus_files], trainer)
tokenizer.post_processor = TemplateProcessing(
single="[BOS] $A [EOS]",
pair="[BOS] $A [SEP] $B [EOS]",
special_tokens=[
("[BOS]", tokenizer.token_to_id("[BOS]")),
("[EOS]", tokenizer.token_to_id("[EOS]")),
("[SEP]", tokenizer.token_to_id("[SEP]")),
],
)
output = Path(output_path)
output.parent.mkdir(parents=True, exist_ok=True)
tokenizer.save(str(output))
logger.info("Tokenizer saved to %s", output)
def load_tokenizer(path: str | Path) -> "DeltaTokenizer":
"""Load a DeltaTokenizer from disk."""
return DeltaTokenizer(path)
class DeltaTokenizer:
"""Thin wrapper around HuggingFace tokenizers.Tokenizer."""
def __init__(self, path: str | Path) -> None:
self.path = Path(path)
self.tokenizer = Tokenizer.from_file(str(self.path))
self.pad_token_id = self.tokenizer.token_to_id("[PAD]")
self.unk_token_id = self.tokenizer.token_to_id("[UNK]")
self.bos_token_id = self.tokenizer.token_to_id("[BOS]")
self.eos_token_id = self.tokenizer.token_to_id("[EOS]")
self.sys_token_id = self.tokenizer.token_to_id("[SYS]")
self.usr_token_id = self.tokenizer.token_to_id("[USR]")
self.ass_token_id = self.tokenizer.token_to_id("[ASS]")
self.sep_token_id = self.tokenizer.token_to_id("[SEP]")
self.special_tokens = set(SPECIAL_TOKENS)
@property
def chat_stop_token_ids(self) -> set[int]:
"""Token ids that should end an assistant completion."""
return {
token_id
for token_id in (
self.eos_token_id,
self.sep_token_id,
self.sys_token_id,
self.usr_token_id,
self.ass_token_id,
self.bos_token_id,
self.pad_token_id,
)
if token_id is not None
}
def encode(self, text: str, add_special_tokens: bool = True) -> list[int]:
"""Encode a string into token ids."""
return self.tokenizer.encode(text, add_special_tokens=add_special_tokens).ids
def decode(self, ids: list[int], skip_special_tokens: bool = True) -> str:
"""Decode token ids into text."""
return self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)
def batch_encode(self, texts: list[str], add_special_tokens: bool = True) -> list[list[int]]:
"""Encode a batch of strings."""
return [encoding.ids for encoding in self.tokenizer.encode_batch(texts, add_special_tokens=add_special_tokens)]
def format_chat(self, messages: list[dict[str, Any]], persona: str | None = None) -> str:
"""Format a multi-turn conversation for Delta.
Args:
messages: Conversation turns with role and content.
persona: Optional system prompt.
Returns:
Prompt text ending with an assistant tag for continuation.
"""
system = persona or DEFAULT_SYSTEM_PROMPT
parts = [f"[SYS] {system} [SEP]"]
for message in messages:
role = str(message.get("role", "")).lower()
content = str(message.get("content", "")).strip()
if role == "user":
parts.append(f"[USR] {content} [SEP]")
elif role == "assistant":
parts.append(f"[ASS] {content} [SEP]")
elif role == "system":
parts[0] = f"[SYS] {content} [SEP]"
if not parts[-1].startswith("[ASS]"):
parts.append("[ASS]")
else:
parts.append("[USR] [SEP]")
parts.append("[ASS]")
return "\n".join(parts)
|