refactor(khaosz): 重构项目结构

khaosz/data/__init__.py

@@ -0,0 +1,30 @@
+from khaosz.data.data_util import (
+    BaseDataset, 
+    SeqDataset, 
+    DpoDataset, 
+    SftDataset, 
+    PpoDataset, 
+    MutiSegmentFetcher,
+    ResumeableRandomSampler,
+    DatasetLoader,
+    load_pkl_files, 
+    build_attention_mask, 
+    build_loss_mask
+)
+
+from khaosz.data.tokenizer import BpeTokenizer
+
+__all__ = [
+    "BaseDataset",
+    "SeqDataset",
+    "DpoDataset",
+    "SftDataset",
+    "PpoDataset",
+    "MutiSegmentFetcher",
+    "ResumeableRandomSampler",
+    "DatasetLoader",
+    "load_pkl_files",
+    "build_attention_mask",
+    "build_loss_mask",
+    "BpeTokenizer"
+]

khaosz/data/data_util.py

@@ -0,0 +1,301 @@
+import torch
+import bisect
+import pickle as pkl
+from abc import ABC, abstractmethod
+from torch import Tensor
+from torch.utils.data import Dataset, Sampler
+from typing import Callable, List, Dict, Literal, Union
+
+MutiSeg = Dict[str, List[Tensor]]
+Seg = Dict[str, Tensor]
+
+def load_pkl_files(paths: List[str]):
+    segments: MutiSeg = {}
+    total_samples = 0
+
+    for path in paths:
+        with open(path, "rb") as f:
+            pkl_file: Seg = pkl.load(f)
+        for key, value in pkl_file.items():
+            if key not in segments:
+                segments[key] = []
+            segments[key].append(value)
+        first_key = list(pkl_file.keys())[0]
+        total_samples += pkl_file[first_key].numel()
+    
+    return segments, total_samples
+
+def build_attention_mask(input_ids: Tensor, user_token_id: int, multi_turn: bool) -> Tensor:
+    seq_len = input_ids.size(0)
+    turn_id = input_ids.eq(user_token_id).cumsum(dim=-1)
+    
+    iq = turn_id.view(seq_len, 1)
+    ik = turn_id.view(1, seq_len)
+    
+    # fix the causual attention mask(iq >= ik condition)
+    seq_mask = (iq >= ik) if multi_turn else (iq == ik)
+    attention_mask = torch.tril(seq_mask)
+    
+    # fix the shape  (bsz, 1, seq_len, seq_len) unsqueeze for broadcast
+    return attention_mask.unsqueeze(0)
+
+def build_loss_mask(input_ids: Tensor, bos_token_id: int, eos_token_id: int) -> Tensor:
+    token_markers = torch.zeros_like(input_ids, dtype=torch.int8)
+    
+    is_bos_token = input_ids.eq(bos_token_id)
+    is_eos_token = input_ids.eq(eos_token_id)
+    
+    # fix the eos_token_id bug(change target_ids to input_ids)
+    token_markers[is_bos_token] = 1
+    token_markers[is_eos_token] = -1 
+    
+    cumulative_markers = torch.cumsum(token_markers, dim=-1)
+    min_cumulative = cumulative_markers.min(dim=-1, keepdim=True).values
+    loss_mask = cumulative_markers - min_cumulative
+    
+    return loss_mask.to(dtype=torch.bool)
+
+
+class BaseSegmentFetcher:
+    def __init__(self, segments: List[Tensor]):
+        self.segments = segments
+        self.cum_lengths = []
+        total = 0
+        for seg in segments:
+            total += len(seg)
+            self.cum_lengths.append(total)
+        self.total_length = total if segments else 0
+
+    def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+        if not (0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length):
+            raise ValueError("begin_idx or end_idx out of bounds")
+        if begin_idx >= end_idx:
+            return torch.tensor([], dtype=torch.long)
+        
+        # fix the range index bug
+        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx)
+        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx)
+
+        result_segments = []
+
+        for i in range(seg_start_idx, seg_end_idx + 1):
+            prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
+            start = max(begin_idx - prev_cum, 0)
+            end = min(end_idx - prev_cum, len(self.segments[i]))
+            result_segments.append(self.segments[i][start:end])
+
+        return torch.cat(result_segments, dim=0)
+    
+
+class MutiSegmentFetcher:
+    def __init__(self, muti_segments: MutiSeg):
+        self.muti_keys = list(muti_segments.keys())
+        self.muti_fetchers = {
+            key: BaseSegmentFetcher(segments)
+            for key, segments in muti_segments.items()
+        }
+        
+    def key_fetch(self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]) -> Union[Tensor, Seg]:
+        fetch_dict = {} 
+        keys = [keys] if isinstance(keys, str) else keys
+        
+        for key in keys:
+            fetcher = self.muti_fetchers[key]
+            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
+            fetch_dict[key] = fetch_tensor
+
+        return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
+    
+    def fetch_data(self, begin_idx: int, end_idx: int) -> Union[Tensor, Seg]:
+        return self.key_fetch(begin_idx, end_idx, self.muti_keys)
+
+
+class BaseDataset(Dataset, ABC):
+    def __init__(self, chunk_size: int):
+        super().__init__()
+        self.segments: MutiSeg = {}
+        self.chunk_size = chunk_size
+        self.total_samples = 0
+
+    def save(self, save_path: str):
+        keys = list(self.segments.keys())
+        if not keys:
+            return
+        
+        first_item = self.segments[keys[0]]
+        segment_size = len(first_item)
+        
+        for i in range(segment_size):
+            formated_segment = {key: self.segments[key][i] for key in keys}
+            pkl.dump(formated_segment, open(f"{save_path}_{i}.pkl", "wb"))
+                
+    
+    def load(self, load_path: Union[str, List[str]]):
+        paths = [load_path] if isinstance(load_path, str) else load_path
+        self.segments, self.total_samples = load_pkl_files(paths)
+        self.fetcher = MutiSegmentFetcher(self.segments)
+        
+    @abstractmethod
+    def __getitem__(self, index: int) -> Dict[str, Tensor]:
+        raise NotImplementedError
+        
+    def __len__(self) -> int:
+        assert self.total_samples // self.chunk_size > 0
+        return self.total_samples // self.chunk_size
+    
+
+class SeqDataset(BaseDataset):
+    def __init__(
+        self, 
+        chunk_size, 
+    ):
+        super().__init__(chunk_size)
+        self.fetcher = MutiSegmentFetcher(self.segments)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, "sequence")
+    
+    def __getitem__(self, index):
+        # fix the range index bug
+        begin_idx = min(index * self.chunk_size, self.total_samples - self.chunk_size - 1)
+        end_idx = begin_idx + self.chunk_size
+        
+        x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long)
+        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long)
+        
+        return {"input_ids": x, "target_ids": y}
+    
+    
+class SftDataset(BaseDataset):
+    def __init__(
+        self, 
+        chunk_size, 
+        bos_token_id, 
+        eos_token_id, 
+        user_token_id, 
+        multi_turn=False, 
+    ):
+        super().__init__(chunk_size)
+        self.fetcher = MutiSegmentFetcher(self.segments)
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.user_token_id = user_token_id
+        self.multi_turn = multi_turn
+    
+    def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, "sequence")
+    
+    def __getitem__(self, index):
+        begin_idx = min(index * self.chunk_size, self.total_samples - self.chunk_size - 1)
+        end_idx = begin_idx + self.chunk_size
+        
+        x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long)
+        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long)
+        
+        # fix the eos_token_id bug(change target_ids to input_ids)
+        loss_mask = build_loss_mask(x, self.bos_token_id, self.eos_token_id)
+        attn_mask = build_attention_mask(x, self.user_token_id, self.multi_turn)
+        
+        return {"input_ids": x, "target_ids": y, "loss_mask": loss_mask, "attn_mask": attn_mask}
+
+
+class DpoDataset(BaseDataset):
+    def __init__(self, chunk_size: int):
+        super().__init__(chunk_size)
+        self.fetcher = MutiSegmentFetcher(self.segments)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+
+    def __getitem__(self, index: int):
+        begin_idx = min(index * self.chunk_size, self.total_samples - self.chunk_size - 1)
+        end_idx = begin_idx + self.chunk_size
+        
+        chosen = self._fetch_data(begin_idx, end_idx, "chosen").to(dtype=torch.long)
+        rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
+        chosen_mask = self._fetch_data(begin_idx, end_idx, "chosen_mask").to(dtype=torch.bool)
+        rejected_mask = self._fetch_data(begin_idx, end_idx, "rejected_mask").to(dtype=torch.bool)
+
+        return {"chosen": chosen, "rejected": rejected, "chosen_mask": chosen_mask, "rejected_mask": rejected_mask}
+
+
+class PpoDataset(BaseDataset):
+    def __init__(self, chunk_size: int):
+        super().__init__(chunk_size)
+        self.fetcher = MutiSegmentFetcher(self.segments)
+
+    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+    
+    def __getitem__(self, index: int) -> Dict[str, Tensor]:
+        begin_idx = min(index * self.chunk_size, self.total_samples - self.chunk_size - 1)
+        end_idx = begin_idx + self.chunk_size
+
+        input_ids =  self._fetch_data(begin_idx, end_idx, "input_ids"),
+        actions = self._fetch_data(begin_idx, end_idx, "actions"),
+        logprobs = self._fetch_data(begin_idx, end_idx, "logprobs"),
+        rewards =  self._fetch_data(begin_idx, end_idx, "rewards")
+        
+        return {"input_ids": input_ids, "actions": actions, "logprobs": logprobs, "rewards": rewards}
+    
+
+class DatasetLoader:
+    @staticmethod       
+    def load(
+        train_type: Literal["seq", "sft", "dpo"],
+        load_path: Union[str, List[str]],
+        max_len: int, 
+        **kwargs
+        ) -> BaseDataset:
+        
+        dataset_router: Dict[str, Callable[[int], BaseDataset]] = {
+            "seq": lambda max_len: SeqDataset(max_len),
+            "sft": lambda max_len: SftDataset(
+                max_len, 
+                bos_token_id=kwargs.get("bos_token_id"),
+                eos_token_id=kwargs.get("eos_token_id"),
+                user_token_id=kwargs.get("user_token_id"),
+                multi_turn=kwargs.get("multi_turn")
+            ),
+            "dpo": lambda max_len: DpoDataset(max_len),
+        }
+        dataset = dataset_router[train_type](max_len)
+        dataset.load(load_path)
+        
+        return dataset
+
+
+class ResumeableRandomSampler(Sampler[int]):
+    def __init__(self, data_source, start_epoch=0, start_iter=0, seed=42):
+        self.num_samples = len(data_source)
+        self.epoch = start_epoch
+        self.iter = start_iter
+        
+        generator = torch.Generator()
+        generator.manual_seed(seed)
+        
+        self.generator = generator
+        self._indices = None 
+    
+    def _get_indices(self):
+        for _ in range(self.epoch):
+            _ = torch.randperm(self.num_samples, generator=self.generator)
+        
+        current_epoch_indices = torch.randperm(self.num_samples, generator=self.generator).tolist()
+        self._indices = current_epoch_indices[self.iter % self.num_samples:]
+    
+    def __iter__(self):
+        if self._indices is None:
+            self._get_indices()
+        
+        for i in self._indices:
+            self.iter += 1
+            yield i
+        
+        self.epoch += 1
+        self._indices = None
+    
+    def __len__(self):
+        if self._indices is None:
+            self._get_indices()
+        return len(self._indices)

khaosz/data/tokenizer.py

@@ -0,0 +1,119 @@
+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
+
+
+class BpeTokenizer:
+    def __init__(self, path=None):
+        self._control_tokens = ["<bos>", "<eos>", "<pad>"]
+        self._special_tokens = ["<|user|>", "<|system|>"]
+        model = BPE()
+        tokenizer = Tokenizer(model)
+        tokenizer.normalizer = normalizers.Sequence([
+            normalizers.NFC()
+        ])
+        tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
+            pre_tokenizers.Punctuation(behavior="isolated"),
+            pre_tokenizers.Metaspace(prepend_scheme="never"),
+            pre_tokenizers.Split(pattern=r"(\d+|[a-zA-Z]+|(?:'s|'t|'re|'ve|'m|'ll|'d))", behavior="isolated"),
+            pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)
+        ])
+        tokenizer.decoder = decoders.Sequence([
+            decoders.ByteLevel(),
+            decoders.Metaspace(prepend_scheme="never")
+        ])
+        tokenizer.post_processor = processors.Sequence([
+            processors.ByteLevel(trim_offsets=False)
+        ])
+        self._tokenizer = tokenizer
+        
+        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) -> tuple:
+        assert reserved_token_size > len(self._special_tokens)
+        reserved_tokens = [f"<|rsv{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 // 4,
+            max_token_length=18,
+            special_tokens=self._control_tokens,
+            show_progress=True,
+            initial_alphabet=alphabet,
+        )
+        
+        return trainer, detail_vocab_size, reserved_tokens
+
+    def train(self, files, vocab_size, min_freq, reserved_token_size=100):
+        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._special_tokens + reserved_tokens)
+            
+    def train_from_iterator(self, iterator, vocab_size, min_freq, reserved_token_size=100):
+        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._special_tokens + reserved_tokens)
+            
+    def save(self, path):
+        self._tokenizer.save(path)
+        
+    def load(self, path):
+        self._tokenizer = Tokenizer.from_file(path)
+
+    def encode(self, tokens: Union[str, List[str]], out_ids: bool=True, add_special_tokens: bool=False) -> List:
+        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_ids = []
+        for token in self._control_tokens:
+            stop_ids.append(self._tokenizer.token_to_id(token))
+        return stop_ids
+    
+    @property
+    def bos_id(self) -> int:
+        return self._tokenizer.token_to_id("<bos>")
+    
+    @property
+    def eos_id(self) -> int:
+        return self._tokenizer.token_to_id("<eos>")
+    
+    @property
+    def pad_id(self) -> int:
+        return self._tokenizer.token_to_id("<pad>")
+    
+    @property
+    def user_id(self) -> int:
+        return self._tokenizer.token_to_id("<|user|>")
+    
+    @property
+    def system_id(self) -> int:
+        return self._tokenizer.token_to_id("<|system|>")