reafactor: 重构项目
This commit is contained in:
+14
-9
@@ -1,17 +1,17 @@
|
||||
import logging
|
||||
from .tokenizer import BpeTokenizer
|
||||
from .text import TextNormalizer
|
||||
from .packing import SequencePacker
|
||||
from .io import IOHandler
|
||||
from .processors import ProcessorFactory, BaseProcessor
|
||||
from .export import export_dataset
|
||||
from .cache import cache_jsonl
|
||||
from .utils import setup_logging
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.text import TextNormalizer
|
||||
from pipeline.packing import SequencePacker
|
||||
from pipeline.io import IOHandler, export_dataset, cache_jsonl
|
||||
from pipeline.processors import ProcessorFactory, BaseProcessor
|
||||
from pipeline.utils import setup_logging
|
||||
from pipeline.strategies import PromptStrategy, ChatMLStrategy, AlpacaStrategy, StrategyFactory
|
||||
|
||||
# 配置项目级日志记录
|
||||
# Configure project-level logging
|
||||
setup_logging()
|
||||
|
||||
__all__ = [
|
||||
# Core modules
|
||||
'BpeTokenizer',
|
||||
'TextNormalizer',
|
||||
'SequencePacker',
|
||||
@@ -20,4 +20,9 @@ __all__ = [
|
||||
'BaseProcessor',
|
||||
'export_dataset',
|
||||
'cache_jsonl',
|
||||
# Strategy pattern
|
||||
'PromptStrategy',
|
||||
'ChatMLStrategy',
|
||||
'AlpacaStrategy',
|
||||
'StrategyFactory',
|
||||
]
|
||||
|
||||
@@ -1,82 +0,0 @@
|
||||
"""Tokenize JSONL files and pack them into HDF5 storage."""
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from typing import List, Dict
|
||||
from pathlib import Path
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from .processors import BaseProcessor
|
||||
from .packing import SequencePacker
|
||||
from .io import IOHandler
|
||||
from .utils import error_handler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@error_handler()
|
||||
def cache_jsonl(
|
||||
files: List[str],
|
||||
output_dir: str,
|
||||
processor: BaseProcessor,
|
||||
*,
|
||||
pack_size: int = -1,
|
||||
pad_value: int = 1,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Tokenize JSONL files and pack them into HDF5 storage.
|
||||
|
||||
Args:
|
||||
files: List of JSONL file paths
|
||||
output_dir: H5 output directory
|
||||
processor: Initialized Processor instance
|
||||
pack_size: Packing length, <=0 means no packing
|
||||
pad_value: Padding value
|
||||
|
||||
Returns:
|
||||
List of generated H5 file paths
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
output_files: List[str] = []
|
||||
# Cache output_keys to avoid repeated attribute access
|
||||
output_keys = processor.output_keys
|
||||
|
||||
for file_path in files:
|
||||
file_name = Path(file_path).stem
|
||||
|
||||
# Pre-allocate lists for each output key
|
||||
arrows: Dict[str, List] = {key: [] for key in output_keys}
|
||||
|
||||
# Read and process all lines
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
for line_num, line in enumerate(tqdm(f, desc=f"Processing {file_name}", leave=False), start=1):
|
||||
try:
|
||||
result = processor.process(json.loads(line))
|
||||
if result is not None:
|
||||
# Batch append: add each key's tensor to corresponding list
|
||||
for key in output_keys:
|
||||
arrows[key].append(result[key])
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line.")
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning(f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line.")
|
||||
continue
|
||||
|
||||
# Convert lists to tensors once per key
|
||||
if pack_size > 0:
|
||||
output = {}
|
||||
for key in output_keys:
|
||||
packer = SequencePacker(pack_size, pad_value)
|
||||
output[key] = packer.pack(arrows[key])
|
||||
else:
|
||||
# No packing: directly use the arrow tensors
|
||||
output = arrows
|
||||
|
||||
IOHandler.save_h5(output_dir, file_name, output)
|
||||
h5_path = os.path.join(output_dir, f"{file_name}.h5")
|
||||
output_files.append(h5_path)
|
||||
logger.info(f"Saved {h5_path}")
|
||||
|
||||
return output_files
|
||||
@@ -1,63 +0,0 @@
|
||||
"""Export HuggingFace Dataset to JSONL files in chunks."""
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from typing import Callable, Optional, List, Union, Dict, Any
|
||||
|
||||
from datasets import Dataset
|
||||
from .utils import error_handler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@error_handler()
|
||||
def export_dataset(
|
||||
dataset: Dataset,
|
||||
output_dir: str,
|
||||
output_prefix: str,
|
||||
*,
|
||||
chunk_size: int = 1_000_000,
|
||||
max_chunks: Optional[int] = None,
|
||||
process_func: Optional[Callable[[Dict[str, Any]], Union[Dict[str, Any], List[Dict[str, Any]]]]] = None,
|
||||
column: str = "text",
|
||||
) -> List[str]:
|
||||
"""
|
||||
Export HuggingFace Dataset to JSONL files in chunks.
|
||||
|
||||
Args:
|
||||
dataset: HuggingFace Dataset object
|
||||
output_dir: Output directory
|
||||
output_prefix: Output file name prefix, e.g., "chinese-c4-pretrain"
|
||||
chunk_size: Maximum number of samples per file
|
||||
max_chunks: Maximum number of chunks to process (for debugging)
|
||||
process_func: Single sample transformation function (dict) -> dict | list[dict]
|
||||
column: Default text column name (only used when process_func is None)
|
||||
|
||||
Returns:
|
||||
List of generated file paths
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
total = len(dataset)
|
||||
num_chunks = (total + chunk_size - 1) // chunk_size
|
||||
lim = min(max_chunks, num_chunks) if max_chunks else num_chunks
|
||||
|
||||
output_files: List[str] = []
|
||||
for i in range(lim):
|
||||
start = i * chunk_size
|
||||
end = min(start + chunk_size, total)
|
||||
chunk = dataset.select(range(start, end))
|
||||
|
||||
path = os.path.join(output_dir, f"{output_prefix}_chunk_{i}.jsonl")
|
||||
try:
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
for example in chunk:
|
||||
processed = process_func(example) if process_func else {column: example[column]}
|
||||
items = processed if isinstance(processed, list) else [processed]
|
||||
for item in items:
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
output_files.append(path)
|
||||
logger.info(f"[{i + 1}/{lim}] Saved {path}")
|
||||
except (OSError, IOError) as e:
|
||||
logger.error(f"Failed to write chunk {i} to {path}: {e}")
|
||||
|
||||
return output_files
|
||||
+143
-10
@@ -1,23 +1,36 @@
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Callable
|
||||
"""File, HDF5, JSONL I/O operations."""
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Callable, Union, Any
|
||||
|
||||
import h5py
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from tqdm import tqdm
|
||||
from datasets import Dataset
|
||||
|
||||
from .utils import error_handler
|
||||
from pipeline.utils import error_handler
|
||||
from pipeline.processors import BaseProcessor
|
||||
from pipeline.packing import SequencePacker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IOHandler:
|
||||
"""File and HDF5 read/write operations."""
|
||||
|
||||
@staticmethod
|
||||
def fetch_files(directory: str) -> List[str]:
|
||||
return [
|
||||
def fetch_files(directory: str, suffix: Optional[str] = None) -> List[str]:
|
||||
files = [
|
||||
os.path.join(root, f)
|
||||
for root, _, files in os.walk(directory)
|
||||
for f in files
|
||||
]
|
||||
if suffix:
|
||||
files = [f for f in files if f.endswith(suffix)]
|
||||
return sorted(files)
|
||||
|
||||
@staticmethod
|
||||
def fetch_folders(root_dir: str, filter_func: Optional[Callable[[str], bool]] = None) -> List[str]:
|
||||
@@ -34,7 +47,7 @@ class IOHandler:
|
||||
def save_h5(output_dir: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
full_path = os.path.join(output_dir, f"{file_name}.h5")
|
||||
|
||||
|
||||
with h5py.File(full_path, 'w') as f:
|
||||
for key, tensors in tensor_group.items():
|
||||
grp = f.create_group(key)
|
||||
@@ -43,12 +56,12 @@ class IOHandler:
|
||||
|
||||
@staticmethod
|
||||
@error_handler()
|
||||
def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
|
||||
def load_h5(file_path: str, share_memory: bool = True) -> Dict[str, List[Tensor]]:
|
||||
tensor_group: Dict[str, List[Tensor]] = {}
|
||||
|
||||
root_path = Path(file_path)
|
||||
h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
|
||||
|
||||
|
||||
for h5_file in h5_files:
|
||||
with h5py.File(h5_file, 'r') as f:
|
||||
for key in f.keys():
|
||||
@@ -60,9 +73,129 @@ class IOHandler:
|
||||
if share_memory:
|
||||
tensor = tensor.share_memory_()
|
||||
dsets.append(tensor)
|
||||
|
||||
|
||||
if tensor_group.get(key) is None:
|
||||
tensor_group[key] = []
|
||||
tensor_group[key].extend(dsets)
|
||||
|
||||
return tensor_group
|
||||
return tensor_group
|
||||
|
||||
|
||||
# ── Stage 1: Export HuggingFace Dataset to JSONL ──────────────────────────
|
||||
|
||||
|
||||
@error_handler()
|
||||
def export_dataset(
|
||||
dataset: Dataset,
|
||||
output_dir: str,
|
||||
output_prefix: str,
|
||||
*,
|
||||
chunk_size: int = 1_000_000,
|
||||
max_chunks: Optional[int] = None,
|
||||
process_func: Optional[Callable[[Dict[str, Any]], Union[Dict[str, Any], List[Dict[str, Any]]]]] = None,
|
||||
column: str = "text",
|
||||
) -> List[str]:
|
||||
"""
|
||||
Export HuggingFace Dataset to JSONL files in chunks.
|
||||
|
||||
Args:
|
||||
dataset: HuggingFace Dataset object
|
||||
output_dir: Output directory
|
||||
output_prefix: Output file name prefix, e.g., "chinese-c4-pretrain"
|
||||
chunk_size: Maximum number of samples per file
|
||||
max_chunks: Maximum number of chunks to process (for debugging)
|
||||
process_func: Single sample transformation function (dict) -> dict | list[dict]
|
||||
column: Default text column name (only used when process_func is None)
|
||||
|
||||
Returns:
|
||||
List of generated file paths
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
total = len(dataset)
|
||||
num_chunks = (total + chunk_size - 1) // chunk_size
|
||||
lim = min(max_chunks, num_chunks) if max_chunks else num_chunks
|
||||
|
||||
output_files: List[str] = []
|
||||
for i in range(lim):
|
||||
start = i * chunk_size
|
||||
end = min(start + chunk_size, total)
|
||||
chunk = dataset.select(range(start, end))
|
||||
|
||||
path = os.path.join(output_dir, f"{output_prefix}_chunk_{i}.jsonl")
|
||||
try:
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
for example in chunk:
|
||||
processed = process_func(example) if process_func else {column: example[column]}
|
||||
items = processed if isinstance(processed, list) else [processed]
|
||||
for item in items:
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
output_files.append(path)
|
||||
logger.info(f"[{i + 1}/{lim}] Saved {path}")
|
||||
except (OSError, IOError) as e:
|
||||
logger.error(f"Failed to write chunk {i} to {path}: {e}")
|
||||
|
||||
return output_files
|
||||
|
||||
|
||||
# ── Stage 2: Tokenize JSONL and cache to HDF5 ────────────────────────────
|
||||
|
||||
|
||||
@error_handler()
|
||||
def cache_jsonl(
|
||||
files: List[str],
|
||||
output_dir: str,
|
||||
processor: BaseProcessor,
|
||||
*,
|
||||
pack_size: int = -1,
|
||||
pad_value: int = 1,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Tokenize JSONL files and pack them into HDF5 storage.
|
||||
|
||||
Args:
|
||||
files: List of JSONL file paths
|
||||
output_dir: H5 output directory
|
||||
processor: Initialized Processor instance
|
||||
pack_size: Packing length, <=0 means no packing
|
||||
pad_value: Padding value
|
||||
|
||||
Returns:
|
||||
List of generated H5 file paths
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
output_files: List[str] = []
|
||||
output_keys = processor.output_keys
|
||||
|
||||
for file_path in files:
|
||||
file_name = Path(file_path).stem
|
||||
|
||||
arrows: Dict[str, List] = {key: [] for key in output_keys}
|
||||
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
for line_num, line in enumerate(tqdm(f, desc=f"Processing {file_name}", leave=False), start=1):
|
||||
try:
|
||||
result = processor.process(json.loads(line))
|
||||
if result is not None:
|
||||
for key in output_keys:
|
||||
arrows[key].append(result[key])
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line.")
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning(f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line.")
|
||||
continue
|
||||
|
||||
if pack_size > 0:
|
||||
output = {}
|
||||
for key in output_keys:
|
||||
packer = SequencePacker(pack_size, pad_value)
|
||||
output[key] = packer.pack(arrows[key])
|
||||
else:
|
||||
output = arrows
|
||||
|
||||
IOHandler.save_h5(output_dir, file_name, output)
|
||||
h5_path = os.path.join(output_dir, f"{file_name}.h5")
|
||||
output_files.append(h5_path)
|
||||
logger.info(f"Saved {h5_path}")
|
||||
|
||||
return output_files
|
||||
|
||||
+94
-73
@@ -1,109 +1,130 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from typing import List
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from .utils import error_handler
|
||||
from pipeline.utils import error_handler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SequencePacker:
|
||||
"""
|
||||
Packs variable-length sequences into fixed-size tensors, suitable for
|
||||
concatenating unequal-length training samples into uniform shapes
|
||||
for DataLoader / model training.
|
||||
|
||||
def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = torch.int32):
|
||||
Algorithm (Sorted Greedy Fill, based on First-Fit Decreasing heuristic):
|
||||
|
||||
Input: sequences = [A(len=5), B(len=2), C(len=3)], pack_size = 8
|
||||
|
||||
1. Validate & Normalize
|
||||
- Check 1D dimension, unify dtype, truncate overlong sequences with warning
|
||||
- Result: [(A,5), (B,2), (C,3)]
|
||||
|
||||
2. Sort by length descending (FFD)
|
||||
- Result: [(A,5), (C,3), (B,2)]
|
||||
|
||||
3. Greedy fill: write into a pre-allocated buffer sequentially, flush when full
|
||||
- Write A(5) -> buffer = [A A A A A _ _ _], pos=5
|
||||
- Write C(3) -> pos+3=8 <= 8 -> buffer = [A A A A A C C C], pos=8
|
||||
- Buffer full -> flush as package[0], reset buffer & pos=0
|
||||
- Write B(2) -> buffer = [B B _ _ _ _ _ _], pos=2
|
||||
- Loop ends -> flush tail -> package[1] = [B B 0 0 0 0 0 0]
|
||||
|
||||
Output: [package[0], package[1]]
|
||||
|
||||
Cross-group consistency:
|
||||
When packing different key groups (e.g. sequences and loss_masks)
|
||||
with separate pack() calls, tensors at the same index always have
|
||||
identical lengths, so the descending sort produces the exact same
|
||||
ordering. Element-level correspondence across groups is preserved.
|
||||
|
||||
Performance:
|
||||
- Pre-allocated buffer reused via fill_() to avoid repeated tensor creation
|
||||
- Attributes cached as local variables inside the loop to reduce lookup overhead
|
||||
"""
|
||||
|
||||
def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = None):
|
||||
self.pack_size = pack_size
|
||||
self.pad_value = pad_value
|
||||
self.dtype = dtype
|
||||
# Pre-allocate buffer for better performance
|
||||
self._buffer: Optional[Tensor] = None
|
||||
self._reset()
|
||||
|
||||
def _reset(self) -> None:
|
||||
"""Reset internal state for instance reuse."""
|
||||
# Reuse buffer instead of creating new tensors
|
||||
if self._buffer is None or self._buffer.shape[0] != self.pack_size:
|
||||
self._buffer = torch.full(
|
||||
(self.pack_size,), self.pad_value, dtype=self.dtype
|
||||
)
|
||||
else:
|
||||
self._buffer.fill_(self.pad_value)
|
||||
self._current_pos = 0
|
||||
self.dtype = dtype # None = follow input dtype
|
||||
self._buffer: Tensor | None = None
|
||||
self._pos = 0
|
||||
self._packages: List[Tensor] = []
|
||||
# Backward compatibility: maintain _current_pack reference
|
||||
self._current_pack = self._buffer
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset packer state for instance reuse, unlocking dtype."""
|
||||
self.dtype = None
|
||||
self._buffer = None
|
||||
self._pos = 0
|
||||
self._packages = []
|
||||
|
||||
@error_handler()
|
||||
def pack(self, sequences: List[Tensor]) -> List[Tensor]:
|
||||
"""
|
||||
Pack sequences into fixed-size packages.
|
||||
Pack sequences into fixed-size packages using First-Fit Decreasing.
|
||||
|
||||
Sequences are sorted by length descending to minimize wasted padding.
|
||||
All tensor groups (e.g. sequences, loss_masks) with matching per-item
|
||||
lengths produce identical ordering, so cross-group correspondence is preserved.
|
||||
|
||||
Args:
|
||||
sequences: List of input tensors
|
||||
sequences: List of 1D input tensors.
|
||||
|
||||
Returns:
|
||||
List of packed tensors, each with length equal to pack_size
|
||||
List of packed tensors, each with length equal to pack_size.
|
||||
"""
|
||||
# Input validation
|
||||
if not sequences:
|
||||
return []
|
||||
|
||||
# Validate and cache tensor sizes in one pass
|
||||
tensor_sizes = []
|
||||
|
||||
# --- validate & normalize in a single pass ---
|
||||
normalized: list[tuple[Tensor, int]] = []
|
||||
target_dtype = self.dtype if self.dtype is not None else sequences[0].dtype
|
||||
for i, seq in enumerate(sequences):
|
||||
if seq.dim() != 1:
|
||||
raise ValueError(
|
||||
f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}"
|
||||
)
|
||||
tensor_sizes.append(seq.numel())
|
||||
if seq.dtype != self.dtype:
|
||||
logger.warning(
|
||||
f"Input tensor dtype {seq.dtype} does not match packer dtype {self.dtype}, "
|
||||
f"will be converted. This may affect packing efficiency."
|
||||
)
|
||||
if seq.dtype != target_dtype:
|
||||
seq = seq.to(target_dtype)
|
||||
length = seq.numel()
|
||||
if length > self.pack_size:
|
||||
seq = seq[: self.pack_size]
|
||||
length = self.pack_size
|
||||
normalized.append((seq, length))
|
||||
|
||||
# Reset state for new packing
|
||||
# --- reset internal state ---
|
||||
buf = self._buffer
|
||||
if buf is None or buf.dtype != target_dtype:
|
||||
buf = torch.full((self.pack_size,), self.pad_value, dtype=target_dtype)
|
||||
self._buffer = buf
|
||||
buf.fill_(self.pad_value)
|
||||
self._pos = 0
|
||||
self._packages = []
|
||||
self._reset()
|
||||
|
||||
# Combine sequences with their sizes for sorting
|
||||
indexed_seqs = list(zip(sequences, tensor_sizes))
|
||||
# Sort by size descending (First-Fit Decreasing algorithm)
|
||||
indexed_seqs.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
for tensor, tensor_size in indexed_seqs:
|
||||
# Truncate sequences that exceed pack_size
|
||||
if tensor_size > self.pack_size:
|
||||
logger.warning(
|
||||
f"Sequence length {tensor_size} exceeds pack_size {self.pack_size}, truncating"
|
||||
)
|
||||
tensor_size = self.pack_size
|
||||
tensor = tensor[: self.pack_size]
|
||||
# --- sort by length descending (FFD heuristic) ---
|
||||
normalized.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Current package is full, create a new one
|
||||
if self._current_pos + tensor_size > self.pack_size:
|
||||
# Finish current package (pad to pack_size)
|
||||
package = self._buffer.clone()
|
||||
self._packages.append(package)
|
||||
# Reset buffer for reuse
|
||||
self._buffer.fill_(self.pad_value)
|
||||
self._current_pos = 0
|
||||
# --- greedy fill ---
|
||||
buf = self._buffer
|
||||
pos = self._pos
|
||||
packages = self._packages
|
||||
pack_size = self.pack_size
|
||||
pad_value = self.pad_value
|
||||
|
||||
# Place tensor in current package
|
||||
self._buffer[self._current_pos : self._current_pos + tensor_size] = tensor
|
||||
self._current_pos += tensor_size
|
||||
for tensor, length in normalized:
|
||||
if pos + length > pack_size:
|
||||
# flush current package
|
||||
packages.append(buf.clone())
|
||||
buf.fill_(pad_value)
|
||||
pos = 0
|
||||
buf[pos : pos + length] = tensor
|
||||
pos += length
|
||||
|
||||
# Handle the last package (pad to pack_size)
|
||||
if self._current_pos > 0:
|
||||
package = self._buffer.clone()
|
||||
self._packages.append(package)
|
||||
# flush the last (possibly partial) package
|
||||
if pos > 0:
|
||||
packages.append(buf.clone())
|
||||
|
||||
# Clear buffer and reset state for backward compatibility
|
||||
self._buffer = None
|
||||
self._current_pack = None
|
||||
self._current_pos = 0
|
||||
|
||||
# write back state
|
||||
self._pos = pos
|
||||
return self._packages
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset packer state for reuse. More efficient than creating a new instance."""
|
||||
self._reset()
|
||||
@@ -1,122 +0,0 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List, Any
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from .tokenizer import BpeTokenizer
|
||||
|
||||
|
||||
class BaseProcessor(ABC):
|
||||
"""Abstract base class for processors."""
|
||||
|
||||
@abstractmethod
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def output_keys(self) -> List[str]:
|
||||
pass
|
||||
|
||||
|
||||
class PreTrainProcessor(BaseProcessor):
|
||||
"""Pre-training data processor."""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
segment = input_dict["text"]
|
||||
tokens = self.tokenizer.encode(f"{segment}<eos>")
|
||||
return {'sequence': torch.tensor(tokens, dtype=torch.int32)}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["sequence"]
|
||||
|
||||
|
||||
class SFTProcessor(BaseProcessor):
|
||||
"""Supervised fine-tuning data processor."""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
query = input_dict["query"]
|
||||
response = input_dict["response"]
|
||||
|
||||
q = self.tokenizer.encode(
|
||||
f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
a = self.tokenizer.encode(f"{response}<|im_end|>\n<eos>")
|
||||
|
||||
q_len = len(q)
|
||||
tokens = torch.tensor(q + a, dtype=torch.int32)
|
||||
loss_mask = torch.zeros(q_len + len(a), dtype=torch.bool)
|
||||
loss_mask[q_len:] = True
|
||||
return {"sequence": tokens, "loss_mask": loss_mask}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["sequence", "loss_mask"]
|
||||
|
||||
|
||||
class DPOProcessor(BaseProcessor):
|
||||
"""DPO preference learning data processor."""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
query = input_dict["query"]
|
||||
chosen_response = input_dict["chosen"]
|
||||
rejected_response = input_dict["rejected"]
|
||||
|
||||
q = self.tokenizer.encode(
|
||||
f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
chosen = self.tokenizer.encode(f"{chosen_response}<|im_end|>\n<eos>")
|
||||
q_len = len(q)
|
||||
chosen_len = len(chosen)
|
||||
|
||||
chosen_tokens = torch.tensor(q + chosen, dtype=torch.int32)
|
||||
chosen_mask = torch.zeros(q_len + chosen_len, dtype=torch.bool)
|
||||
chosen_mask[q_len:] = True
|
||||
|
||||
rejected = self.tokenizer.encode(f"{rejected_response}<|im_end|>\n<eos>")
|
||||
rejected_len = len(rejected)
|
||||
|
||||
rejected_tokens = torch.tensor(q + rejected, dtype=torch.int32)
|
||||
rejected_mask = torch.zeros(q_len + rejected_len, dtype=torch.bool)
|
||||
rejected_mask[q_len:] = True
|
||||
|
||||
return {
|
||||
"chosen": chosen_tokens,
|
||||
"chosen_mask": chosen_mask,
|
||||
"rejected": rejected_tokens,
|
||||
"rejected_mask": rejected_mask,
|
||||
}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
|
||||
|
||||
class ProcessorFactory:
|
||||
"""Processor factory."""
|
||||
|
||||
_processors = {
|
||||
"pt": PreTrainProcessor,
|
||||
"sft": SFTProcessor,
|
||||
"dpo": DPOProcessor,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> BaseProcessor:
|
||||
if processor_type not in cls._processors:
|
||||
raise ValueError(f"Invalid processor type: {processor_type}")
|
||||
return cls._processors[processor_type](tokenizer)
|
||||
|
||||
@classmethod
|
||||
def register(cls, processor_type: str, processor_class: type):
|
||||
cls._processors[processor_type] = processor_class
|
||||
@@ -0,0 +1,18 @@
|
||||
"""Data processors with factory pattern.
|
||||
|
||||
Processor classes are registered at definition time via decorators and
|
||||
can be created through :class:`ProcessorFactory`.
|
||||
"""
|
||||
from pipeline.processors.base import BaseProcessor
|
||||
from pipeline.processors.factory import ProcessorFactory
|
||||
from pipeline.processors.pretrain import PreTrainProcessor
|
||||
from pipeline.processors.sft import SFTProcessor
|
||||
from pipeline.processors.dpo import DPOProcessor
|
||||
|
||||
__all__ = [
|
||||
"BaseProcessor",
|
||||
"ProcessorFactory",
|
||||
"PreTrainProcessor",
|
||||
"SFTProcessor",
|
||||
"DPOProcessor",
|
||||
]
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Processor base class and shared utilities."""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List, Any, Tuple
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
def _encode_with_mask(
|
||||
prompt_tokens: List[int],
|
||||
response_tokens: List[int],
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
"""Concatenate token lists and build loss mask (prompt=False, response=True)."""
|
||||
q_len = len(prompt_tokens)
|
||||
combined = torch.tensor(prompt_tokens + response_tokens, dtype=torch.int32)
|
||||
mask = torch.zeros(q_len + len(response_tokens), dtype=torch.bool)
|
||||
mask[q_len:] = True
|
||||
return combined, mask
|
||||
|
||||
|
||||
class BaseProcessor(ABC):
|
||||
"""Abstract base class for processors."""
|
||||
|
||||
@abstractmethod
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def output_keys(self) -> List[str]:
|
||||
pass
|
||||
@@ -0,0 +1,51 @@
|
||||
"""DPO preference learning data processor."""
|
||||
from typing import Dict, List, Any, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies import PromptStrategy, ChatMLStrategy
|
||||
from pipeline.processors.base import BaseProcessor, _encode_with_mask
|
||||
from pipeline.processors.factory import ProcessorFactory
|
||||
|
||||
|
||||
@ProcessorFactory.register("dpo")
|
||||
class DPOProcessor(BaseProcessor):
|
||||
"""DPO preference learning data processor.
|
||||
|
||||
Supports custom prompt strategy via constructor parameter.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: BpeTokenizer,
|
||||
strategy: Optional[PromptStrategy] = None,
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.strategy = strategy or ChatMLStrategy(tokenizer)
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
query_tokens = self.tokenizer.encode(input_dict["query"])
|
||||
chosen_tokens = self.tokenizer.encode(input_dict["chosen"])
|
||||
rejected_tokens = self.tokenizer.encode(input_dict["rejected"])
|
||||
|
||||
prompt = self.strategy.assemble_prompt(query_tokens)
|
||||
|
||||
chosen_t, chosen_m = _encode_with_mask(
|
||||
prompt, self.strategy.assemble_response(chosen_tokens)
|
||||
)
|
||||
rejected_t, rejected_m = _encode_with_mask(
|
||||
prompt, self.strategy.assemble_response(rejected_tokens)
|
||||
)
|
||||
|
||||
return {
|
||||
"chosen": chosen_t,
|
||||
"chosen_mask": chosen_m,
|
||||
"rejected": rejected_t,
|
||||
"rejected_mask": rejected_m,
|
||||
}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
@@ -0,0 +1,122 @@
|
||||
"""Factory for creating and registering processors."""
|
||||
from typing import Dict, List, Any, Optional, Type
|
||||
|
||||
from pipeline.processors.base import BaseProcessor
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies import PromptStrategy, StrategyFactory
|
||||
|
||||
|
||||
class ProcessorFactory:
|
||||
"""Registry and factory for BaseProcessor implementations.
|
||||
|
||||
Supports decorator-based registration for extensible processor types.
|
||||
|
||||
Example usage::
|
||||
|
||||
@ProcessorFactory.register("custom")
|
||||
class CustomProcessor(BaseProcessor):
|
||||
...
|
||||
|
||||
processor = ProcessorFactory.create(optimizer, "custom", **kwargs)
|
||||
"""
|
||||
|
||||
PROCESSOR_MAP: Dict[str, Type[BaseProcessor]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str):
|
||||
"""Decorator to register a new processor class.
|
||||
|
||||
Args:
|
||||
name: Registration name for the processor.
|
||||
|
||||
Returns:
|
||||
Decorator function that registers the processor class.
|
||||
"""
|
||||
|
||||
def decorator(processor_cls: Type[BaseProcessor]) -> Type[BaseProcessor]:
|
||||
if not issubclass(processor_cls, BaseProcessor):
|
||||
raise TypeError(
|
||||
f"{processor_cls.__name__} must inherit from BaseProcessor"
|
||||
)
|
||||
cls.PROCESSOR_MAP[name] = processor_cls
|
||||
return processor_cls
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> BaseProcessor:
|
||||
"""Create a processor by type name (uses default ChatMLStrategy for SFT/DPO).
|
||||
|
||||
Args:
|
||||
processor_type: Registered processor name (e.g. ``"pt"``, ``"sft"``, ``"dpo"``).
|
||||
tokenizer: Tokenizer instance.
|
||||
|
||||
Returns:
|
||||
Processor instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If processor_type is not registered.
|
||||
"""
|
||||
if processor_type not in cls.PROCESSOR_MAP:
|
||||
raise ValueError(
|
||||
f"Unknown processor type: '{processor_type}'. "
|
||||
f"Supported types: {sorted(cls.PROCESSOR_MAP.keys())}"
|
||||
)
|
||||
return cls.PROCESSOR_MAP[processor_type](tokenizer)
|
||||
|
||||
@classmethod
|
||||
def create_with_strategy(
|
||||
cls,
|
||||
processor_type: str,
|
||||
tokenizer: BpeTokenizer,
|
||||
strategy: PromptStrategy,
|
||||
) -> BaseProcessor:
|
||||
"""Create a processor with a custom strategy.
|
||||
|
||||
Only SFT and DPO processors accept a strategy; PreTrain ignores it.
|
||||
|
||||
Args:
|
||||
processor_type: Registered processor name.
|
||||
tokenizer: Tokenizer instance.
|
||||
strategy: Prompt strategy instance.
|
||||
|
||||
Returns:
|
||||
Processor instance configured with strategy.
|
||||
"""
|
||||
if processor_type not in cls.PROCESSOR_MAP:
|
||||
raise ValueError(
|
||||
f"Unknown processor type: '{processor_type}'. "
|
||||
f"Supported types: {sorted(cls.PROCESSOR_MAP.keys())}"
|
||||
)
|
||||
|
||||
processor_cls = cls.PROCESSOR_MAP[processor_type]
|
||||
if processor_type == "pt":
|
||||
return processor_cls(tokenizer)
|
||||
return processor_cls(tokenizer, strategy=strategy)
|
||||
|
||||
@classmethod
|
||||
def create_with_strategy_name(
|
||||
cls,
|
||||
processor_type: str,
|
||||
tokenizer: BpeTokenizer,
|
||||
strategy_name: str,
|
||||
**strategy_kwargs,
|
||||
) -> BaseProcessor:
|
||||
"""Create a processor with a strategy selected by name.
|
||||
|
||||
Args:
|
||||
processor_type: Registered processor name.
|
||||
tokenizer: Tokenizer instance.
|
||||
strategy_name: Registered strategy name (``"chatml"``, ``"alpaca"``, etc.).
|
||||
**strategy_kwargs: Forwarded to the strategy constructor.
|
||||
|
||||
Returns:
|
||||
Processor instance.
|
||||
"""
|
||||
strategy = StrategyFactory.create(strategy_name, tokenizer, **strategy_kwargs)
|
||||
return cls.create_with_strategy(processor_type, tokenizer, strategy)
|
||||
|
||||
@classmethod
|
||||
def available_types(cls) -> List[str]:
|
||||
"""Return list of registered processor type names."""
|
||||
return list(cls.PROCESSOR_MAP.keys())
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Pre-training data processor."""
|
||||
from typing import Dict, List, Any
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.processors.base import BaseProcessor
|
||||
from pipeline.processors.factory import ProcessorFactory
|
||||
|
||||
|
||||
@ProcessorFactory.register("pt")
|
||||
class PreTrainProcessor(BaseProcessor):
|
||||
"""Pre-training data processor."""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
segment = input_dict["text"]
|
||||
tokens = self.tokenizer.encode(f"{segment}<eos>")
|
||||
return {"sequence": torch.tensor(tokens, dtype=torch.int32)}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["sequence"]
|
||||
@@ -0,0 +1,40 @@
|
||||
"""Supervised fine-tuning data processor."""
|
||||
from typing import Dict, List, Any, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies import PromptStrategy, ChatMLStrategy
|
||||
from pipeline.processors.base import BaseProcessor, _encode_with_mask
|
||||
from pipeline.processors.factory import ProcessorFactory
|
||||
|
||||
|
||||
@ProcessorFactory.register("sft")
|
||||
class SFTProcessor(BaseProcessor):
|
||||
"""Supervised fine-tuning data processor.
|
||||
|
||||
Supports custom prompt strategy via constructor parameter.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: BpeTokenizer,
|
||||
strategy: Optional[PromptStrategy] = None,
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.strategy = strategy or ChatMLStrategy(tokenizer)
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
query_tokens = self.tokenizer.encode(input_dict["query"])
|
||||
response_tokens = self.tokenizer.encode(input_dict["response"])
|
||||
|
||||
prompt = self.strategy.assemble_prompt(query_tokens)
|
||||
response = self.strategy.assemble_response(response_tokens)
|
||||
|
||||
tokens, loss_mask = _encode_with_mask(prompt, response)
|
||||
return {"sequence": tokens, "loss_mask": loss_mask}
|
||||
|
||||
@property
|
||||
def output_keys(self) -> List[str]:
|
||||
return ["sequence", "loss_mask"]
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Strategy pattern for prompt/response format abstraction."""
|
||||
from pipeline.strategies.base import PromptStrategy
|
||||
from pipeline.strategies.factory import StrategyFactory
|
||||
|
||||
# Import strategy implementations to trigger decorator registration
|
||||
from pipeline.strategies.chatml import ChatMLStrategy # noqa: F401
|
||||
from pipeline.strategies.alpaca import AlpacaStrategy # noqa: F401
|
||||
|
||||
__all__ = [
|
||||
"PromptStrategy",
|
||||
"StrategyFactory",
|
||||
"ChatMLStrategy",
|
||||
"AlpacaStrategy",
|
||||
]
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Alpaca format strategy."""
|
||||
from typing import List
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies.base import PromptStrategy
|
||||
from pipeline.strategies.factory import StrategyFactory
|
||||
|
||||
|
||||
@StrategyFactory.register("alpaca")
|
||||
class AlpacaStrategy(PromptStrategy):
|
||||
"""Alpaca format: ``### Instruction: ... \\n\\n### Response: ... <eos>``"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: BpeTokenizer,
|
||||
instruction_start: str = "### Instruction:\n",
|
||||
response_start: str = "### Response:\n",
|
||||
response_suffix: str = "\n<eos>",
|
||||
):
|
||||
super().__init__(tokenizer)
|
||||
self.instruction_start = instruction_start
|
||||
self.response_start = response_start
|
||||
self.response_suffix = response_suffix
|
||||
|
||||
self._instruction_start_ids = self._encode_format(instruction_start)
|
||||
self._separator_ids = self._encode_format("\n\n")
|
||||
self._response_start_ids = self._encode_format(response_start)
|
||||
self._response_suffix_ids = self._encode_format(response_suffix)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "alpaca"
|
||||
|
||||
def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
|
||||
return (
|
||||
self._instruction_start_ids
|
||||
+ query_tokens
|
||||
+ self._separator_ids
|
||||
+ self._response_start_ids
|
||||
)
|
||||
|
||||
def assemble_response(self, response_tokens: List[int]) -> List[int]:
|
||||
return response_tokens + self._response_suffix_ids
|
||||
@@ -0,0 +1,41 @@
|
||||
"""Abstract base class for prompt construction strategies."""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
|
||||
|
||||
class PromptStrategy(ABC):
|
||||
"""Abstract base for prompt/response format strategies.
|
||||
|
||||
Strategies operate at the token level: the Processor tokenizes raw
|
||||
text (query, response, …) and passes token lists to the Strategy,
|
||||
which assembles them with pre-encoded format tokens.
|
||||
"""
|
||||
|
||||
def __init__(self, tokenizer: BpeTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def _encode_format(self, text: str) -> List[int]:
|
||||
"""Encode a format string that may contain special tokens."""
|
||||
return self.tokenizer.encode(text, add_special_tokens=False)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
|
||||
"""Assemble query tokens into a complete prompt with format tokens.
|
||||
|
||||
The prompt includes all tokens up to (and including) the response
|
||||
start marker, e.g. ``<|im_start|>assistant\n``.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def assemble_response(self, response_tokens: List[int]) -> List[int]:
|
||||
"""Wrap response tokens with format tokens (suffix, eos, etc)."""
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(name='{self.name}')"
|
||||
@@ -0,0 +1,41 @@
|
||||
"""ChatML format strategy."""
|
||||
from typing import List
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies.base import PromptStrategy
|
||||
from pipeline.strategies.factory import StrategyFactory
|
||||
|
||||
|
||||
@StrategyFactory.register("chatml")
|
||||
class ChatMLStrategy(PromptStrategy):
|
||||
"""ChatML format: ``<|im_start|>user ... <|im_end|> <|im_start|>assistant ... <|im_end|> <eos>``"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: BpeTokenizer,
|
||||
user_start: str = "<|im_start|>user\n",
|
||||
user_end: str = "<|im_end|>\n",
|
||||
assistant_start: str = "<|im_start|>assistant\n",
|
||||
assistant_end: str = "<|im_end|>\n<eos>",
|
||||
):
|
||||
super().__init__(tokenizer)
|
||||
|
||||
self._user_start_ids = self._encode_format(user_start)
|
||||
self._user_end_ids = self._encode_format(user_end)
|
||||
self._assistant_start_ids = self._encode_format(assistant_start)
|
||||
self._assistant_end_ids = self._encode_format(assistant_end)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "chatml"
|
||||
|
||||
def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
|
||||
return (
|
||||
self._user_start_ids
|
||||
+ query_tokens
|
||||
+ self._user_end_ids
|
||||
+ self._assistant_start_ids
|
||||
)
|
||||
|
||||
def assemble_response(self, response_tokens: List[int]) -> List[int]:
|
||||
return response_tokens + self._assistant_end_ids
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Factory for creating and registering prompt strategies."""
|
||||
from typing import Dict, List, Type
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
from pipeline.strategies.base import PromptStrategy
|
||||
|
||||
|
||||
class StrategyFactory:
|
||||
"""Registry and factory for PromptStrategy implementations.
|
||||
|
||||
Supports decorator-based registration for extensible strategy types.
|
||||
|
||||
Example usage::
|
||||
|
||||
@StrategyFactory.register("custom")
|
||||
class CustomStrategy(PromptStrategy):
|
||||
...
|
||||
|
||||
strategy = StrategyFactory.create("custom", tokenizer, **kwargs)
|
||||
"""
|
||||
|
||||
STRATEGY_MAP: Dict[str, Type[PromptStrategy]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str):
|
||||
"""Decorator to register a new strategy class.
|
||||
|
||||
Args:
|
||||
name: Registration name for the strategy.
|
||||
|
||||
Returns:
|
||||
Decorator function that registers the strategy class.
|
||||
"""
|
||||
|
||||
def decorator(strategy_cls: Type[PromptStrategy]) -> Type[PromptStrategy]:
|
||||
if not issubclass(strategy_cls, PromptStrategy):
|
||||
raise TypeError(
|
||||
f"{strategy_cls.__name__} must inherit from PromptStrategy"
|
||||
)
|
||||
cls.STRATEGY_MAP[name] = strategy_cls
|
||||
return strategy_cls
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def create(cls, name: str, tokenizer: BpeTokenizer, **kwargs) -> PromptStrategy:
|
||||
"""Create a strategy by name.
|
||||
|
||||
Args:
|
||||
name: Registered strategy name (e.g. ``"chatml"``, ``"alpaca"``).
|
||||
tokenizer: Tokenizer instance (required by all strategies).
|
||||
**kwargs: Forwarded to the strategy constructor.
|
||||
|
||||
Returns:
|
||||
Strategy instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If name is not registered.
|
||||
"""
|
||||
if name not in cls.STRATEGY_MAP:
|
||||
raise ValueError(
|
||||
f"Unknown strategy: '{name}'. "
|
||||
f"Supported types: {sorted(cls.STRATEGY_MAP.keys())}"
|
||||
)
|
||||
return cls.STRATEGY_MAP[name](tokenizer=tokenizer, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def available_types(cls) -> List[str]:
|
||||
"""Return list of registered strategy type names."""
|
||||
return list(cls.STRATEGY_MAP.keys())
|
||||
Reference in New Issue
Block a user