Initial commit

khaosz/__init__.py

@@ -0,0 +1,52 @@
+__version__ = "1.2.1"
+__author__ = "ViperEkura"
+
+from khaosz.model import Khaosz
+from khaosz.core.transformer import Transformer, TransformerConfig
+from khaosz.utils.retriever import Retriever
+from khaosz.utils.splitter import (
+    SemanticTextSplitter, 
+    PriorityTextSplitter
+)
+from khaosz.core.tokenizer import BpeTokenizer
+from khaosz.core.parameter import ParameterLoader
+from khaosz.core.generator import (
+    TextGenerator,
+    ChatGenerator, 
+    StreamGenerator, 
+    BatchGenerator, 
+    RetrievalGenerator, 
+    EmbeddingEncoder
+)
+from khaosz.trainer.trainer import Trainer
+from khaosz.trainer.dataset import SeqDataset, SftDataset, DpoDataset, BaseDataset
+
+
+__all__ = [
+    # model
+    "Khaosz",
+    
+    # module
+    "Transformer",
+    "TransformerConfig",
+    "BpeTokenizer",
+    "ParameterLoader",
+    "TextGenerator",
+    "ChatGenerator",
+    "StreamGenerator",
+    "BatchGenerator",
+    "RetrievalGenerator",
+    "EmbeddingEncoder",
+    
+    # trainer
+    "Trainer",
+    "SeqDataset",
+    "SftDataset",
+    "DpoDataset",
+    "BaseDataset",
+    
+    # utils
+    "Retriever",
+    "SemanticTextSplitter",
+    "PriorityTextSplitter",
+]

khaosz/core/__init__.py

@@ -0,0 +1,27 @@
+from khaosz.core.tokenizer import BpeTokenizer
+from khaosz.core.transformer import Transformer, TransformerConfig
+from khaosz.core.parameter import ParameterLoader, ModelParameter, Checkpoint
+from khaosz.core.generator import (
+    TextGenerator,
+    ChatGenerator, 
+    StreamGenerator, 
+    BatchGenerator, 
+    RetrievalGenerator, 
+    EmbeddingEncoder
+)
+
+
+__all__ = [
+    "Transformer",
+    "TransformerConfig",
+    "BpeTokenizer",
+    "ParameterLoader",
+    "ModelParameter",
+    "Checkpoint",
+    "TextGenerator",
+    "ChatGenerator",
+    "StreamGenerator",
+    "BatchGenerator",
+    "RetrievalGenerator",
+    "EmbeddingEncoder"
+]

khaosz/core/generator.py

@@ -0,0 +1,568 @@
+import torch
+from torch import Tensor 
+from typing import List, Tuple, Union, Optional, Generator, Self
+from khaosz.core.parameter import ModelParameter
+
+
+def build_prompt(query: str, history: Optional[List[Tuple[str, str]]] = None) -> str:
+    """ 
+    Build prompt for query and history
+    
+    Args:
+        query(str): query string
+        history(Optional[List[Tuple[str, str]]]): history list of query and response
+        
+    Returns:
+        str: prompt string
+        
+    """
+    prompt_parts = []
+    
+    if history is None:
+        history = []
+    
+    for his_query, his_response in history:
+        prompt_parts.append(f"<|user|> {his_query} <|system|> <bos>{his_response}<eos>")
+        
+    if query is not None:
+        prompt_parts.append(f"<|user|> {query} <|system|> <bos>")
+    
+    return "\n".join(prompt_parts)
+
+def pad_sequence(ids_list: List[List[int]], max_ids_len: int, pad_id: int) -> List[List[int]]:
+    """ 
+    Pad a list of sequences to a fixed length.
+    
+    Args:
+        ids_list (List[List[int]]): A list of sequences.
+        max_ids_len (int): The maximum length of sequences.
+        pad_id (int): The id to pad sequences.
+        
+    Returns:
+        List[List[int]]: A list of padded sequences.
+        
+    """
+    new_ids_list = []
+    for ids in ids_list:
+        pad_len = max_ids_len - len(ids)
+        padded_seq = [pad_id] * pad_len + ids
+        new_ids_list.append(padded_seq)
+    
+    return new_ids_list
+
+def apply_sampling_strategies(
+    logits: Tensor,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+    filter_value: float = -float("inf")
+) -> Tensor:
+    """ 
+    Apply sampling strategies to the logits tensor.
+    
+    Args:
+        logits (Tensor): The logits tensor.
+        temperature (float): The temperature parameter.
+        top_k (int): The top-k parameter.
+        top_p (float): The top-p parameter.
+        filter_value (float, optional): The filter value. Defaults to -float("inf").
+        
+    Returns:
+        Tensor: The sampled logits tensor.
+        
+    """
+    
+    if temperature != 1.0:
+        logits = logits / temperature
+    
+    if top_k > 0:
+        top_k = min(top_k, logits.size(-1))
+        indices_to_remove = logits < torch.topk(logits, top_k, dim=-1)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+    
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+        
+        sorted_indices_to_remove = cumulative_probs > top_p
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+        
+        indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
+        indices_to_remove.scatter_(
+            dim=1,
+            index=sorted_indices,
+            src=sorted_indices_to_remove
+        )
+        
+        logits[indices_to_remove] = filter_value
+    
+    return logits
+
+
+class KVCacheManager:
+    def __init__(
+        self, 
+        num_layers: int, 
+        batch_size: int,
+        max_len: int, 
+        num_heads: int, 
+        head_dim: int, 
+        device: torch.device = "cuda", 
+        dtype: torch.dtype = torch.bfloat16
+    ):
+        self.num_layers = num_layers
+        self.batch_size = batch_size
+        self.max_len = max_len
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.device = device
+        self.dtype = dtype
+        
+        self._kv_cache: List[Tuple[Tensor, Tensor]] = None
+        self._seq_mask: Tensor = None
+        self._initialize()
+
+    def _initialize(self):
+        self._kv_cache = []
+        for _ in range(self.num_layers):
+            k_cache = torch.zeros(
+                (self.batch_size, self.max_len, self.num_heads, self.head_dim),
+                device=self.device, dtype=self.dtype
+            )
+            v_cache = torch.zeros(
+                (self.batch_size, self.max_len, self.num_heads, self.head_dim),
+                device=self.device, dtype=self.dtype
+            )
+            self._kv_cache.append((k_cache, v_cache))
+        
+        self._seq_mask = torch.ones(
+            (self.batch_size, self.max_len),
+            device=self.device, dtype=torch.bool
+        )
+
+    def update(self, active_mask: Tensor):
+        for i in range(self.num_layers):
+            k_cache, v_cache = self._kv_cache[i]
+            new_k_cache, new_v_cache = k_cache[active_mask], v_cache[active_mask]
+            self._kv_cache[i] = (new_k_cache, new_v_cache)
+        
+        self._seq_mask = self._seq_mask[active_mask]
+
+    def reset(self, full_reset=False):
+        if full_reset:
+            self._kv_cache = None
+            self._seq_mask = None
+        else:
+            self._initialize()
+    
+    def set_seq_mask(self, input_ids: Tensor, pad_id: int):
+        batch_size, seq_len = input_ids.shape
+        bool_mask = (input_ids != pad_id)
+        self._seq_mask[: batch_size, : seq_len] = bool_mask
+
+    def get_kvcache(self) -> List[Tuple[Tensor, Tensor]]:
+        return self._kv_cache
+    
+    def get_seq_mask(self) -> Tensor:        
+        return self._seq_mask
+    
+
+class GeneratorCore:
+    def __init__(self, parameter: ModelParameter):
+        self.model = parameter.model
+        self.tokenizer = parameter.tokenizer
+        self.config = parameter.config
+
+    def compute_logits(
+        self,
+        input_ids: Tensor,
+        attn_mask: Optional[Tensor] = None,
+        kv_caches: Optional[List[Tuple[Tensor, Tensor]]] = None,
+        start_pos: int = 0
+    ) -> Tuple[Tensor, int]:
+        with torch.inference_mode():
+            outputs = self.model(input_ids, attn_mask, kv_caches, start_pos)
+            logits = outputs["logits"][:, -1, :]
+            cache_increase = input_ids.size(-1)   
+        
+        return logits, cache_increase
+
+    def to(self, *args, **kargs) -> Self:
+        self.model.to(*args, **kargs)
+        return self
+
+
+class EmbeddingEncoderCore:
+    def __init__(self, parameter: ModelParameter):
+        self.model = parameter.model
+        self.tokenizer = parameter.tokenizer
+        self.config = parameter.config
+    
+    def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
+        with_batch = isinstance(sentence, list)
+        ids = self.tokenizer.encode(sentence)
+        batch_ids = ids if with_batch else [ids]
+        max_model_len = self.config.m_len
+        
+        all_fragments = []
+        fragment_origin_idx = []
+        
+        for i, seq in enumerate(batch_ids):
+            if len(seq) > max_model_len:
+                fragments = [seq[j:j+max_model_len] for j in range(0, len(seq), max_model_len)]
+                all_fragments.extend(fragments)
+                fragment_origin_idx.extend([i] * len(fragments))
+            else:
+                all_fragments.append(seq)
+                fragment_origin_idx.append(i)
+        
+        #if empty fragments
+        if not all_fragments or not ids:
+            return [] if with_batch else torch.tensor([])
+        
+        device = next(self.model.parameters()).device
+        max_len = min(max(len(seq) for seq in all_fragments), max_model_len)
+        
+        padded_ids = []
+        masks = []
+        for seq in all_fragments:
+            pad_len = max_len - len(seq)
+            padded_seq = seq + [self.tokenizer.pad_id] * pad_len
+            mask = [token_id != self.tokenizer.pad_id for token_id in padded_seq]
+            padded_ids.append(padded_seq)
+            masks.append(mask)
+        
+        input_tensor = torch.tensor(padded_ids, device=device, dtype=torch.long)
+        seq_mask = torch.tensor(masks, device=device, dtype=torch.bool)
+        
+        with torch.inference_mode():
+            outputs = self.model(input_tensor, seq_mask)["hidden_states"]
+            # [num_fragments, seq_len, hidden_size]
+            fragment_embs = torch.mul(outputs, seq_mask.unsqueeze(-1))  
+        
+        sentence_embs: List[Tensor] = []
+        for i in range(len(batch_ids)):
+            indices = [idx for idx, orig_idx in enumerate(fragment_origin_idx) if orig_idx == i]
+            if indices is not None:
+                sum_frags = torch.sum(fragment_embs[indices, :, :], dim=1)      # [frags, hidden_size]
+                length = torch.sum(seq_mask[indices, :], dim=1).unsqueeze(1)    # [frags, 1]
+                emb = torch.sum(sum_frags / length, dim=0)                      # [frags, hidden_size]
+                sentence_embs.append(emb.flatten())
+        
+        if with_batch:
+            return [emb.flatten() for emb in sentence_embs]
+        else:
+            return sentence_embs[0].flatten()
+
+    def to(self, *args, **kargs) -> Self:
+        self.model.to(*args, **kargs)
+        return self
+
+
+class TextGenerator(GeneratorCore):
+    def __init__(self, parameter: ModelParameter):
+        super().__init__(parameter)
+    
+    def generate(
+            self, 
+            query: str, 
+            temperature: float,
+            top_k: int,
+            top_p: float,
+        ) -> str:
+        assert temperature >= 0.0
+        assert top_k >= 0
+        assert top_p >= 0.0 and top_p <= 1.0
+        
+        device = next(self.model.parameters()).device
+        cache_manager = KVCacheManager(
+            num_layers=self.config.n_layer, 
+            batch_size=1, 
+            max_len=self.config.m_len, 
+            num_heads=self.config.n_kvhead,
+            head_dim=self.config.n_dim // self.config.n_head,
+            device=device,
+        )
+        
+        ids = self.tokenizer.encode(query)
+        input_ids = torch.tensor([ids], device=device, dtype=torch.long)
+        
+        start_cache_pos = len(ids)
+        cur_cache_pos = 0
+        self.model.eval()
+        
+        while len(ids) < self.config.m_len:
+            kv_caches = cache_manager.get_kvcache()
+            logits, cache_increase = self.compute_logits(
+                input_ids,
+                kv_caches=kv_caches, 
+                start_pos=cur_cache_pos
+            )
+            logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
+            probs = torch.softmax(logits, dim=-1)
+            next_token_id = torch.multinomial(probs, num_samples=1)
+            
+            input_ids = next_token_id
+            ids.append(next_token_id.item())
+            cur_cache_pos += cache_increase
+            
+            if next_token_id.item() in self.tokenizer.stop_ids:
+                break
+        
+        response = self.tokenizer.decode(ids[start_cache_pos:])
+        
+        return response
+
+
+
+class ChatGenerator(GeneratorCore):
+    def __init__(self, parameter: ModelParameter):
+        super().__init__(parameter)
+        
+    def generate(
+            self, 
+            query: str, 
+            history: List[Tuple[str, str]],
+            temperature: float,
+            top_k: int,
+            top_p: float,
+        ) -> str:
+        
+        assert temperature >= 0.0
+        assert top_k >= 0
+        assert top_p >= 0.0 and top_p <= 1.0
+
+        if history is None:
+            history = []
+        
+        device = next(self.model.parameters()).device
+        cache_manager = KVCacheManager(
+            num_layers=self.config.n_layer, 
+            batch_size=1, 
+            max_len=self.config.m_len, 
+            num_heads=self.config.n_kvhead,
+            head_dim=self.config.n_dim // self.config.n_head,
+            device=device,
+        )
+        ids = self.tokenizer.encode(build_prompt(query, history))
+        input_ids = torch.tensor([ids], device=device, dtype=torch.long)
+        cpy_history = history.copy()
+        
+        start_cache_pos = len(ids)
+        cur_cache_pos = 0
+        self.model.eval()
+        
+        
+        while len(ids) < self.config.m_len:
+            kv_caches = cache_manager.get_kvcache()
+            logits, cache_increase = self.compute_logits(
+                input_ids,
+                kv_caches=kv_caches, 
+                start_pos=cur_cache_pos
+            )
+            logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
+            probs = torch.softmax(logits, dim=-1)
+            next_token_id = torch.multinomial(probs, num_samples=1)
+            
+            input_ids = next_token_id
+            ids.append(next_token_id.item())
+            cur_cache_pos += cache_increase
+            
+            if next_token_id.item() in self.tokenizer.stop_ids:
+                break
+        
+        response = self.tokenizer.decode(ids[start_cache_pos:])
+        cpy_history.append((query, response))
+        
+        return response, cpy_history
+    
+
+class StreamGenerator(GeneratorCore):
+    def __init__(self, parameter: ModelParameter):
+        super().__init__(parameter)
+    
+    def generate(
+            self, 
+            query: str, 
+            history: List[Tuple[str, str]],
+            temperature: float,
+            top_k: int,
+            top_p: float,
+        ) -> Generator[Tuple[str, List[Tuple[str, str]]], None, None]:
+        
+        assert temperature >= 0.0
+        assert top_k >= 0
+        assert top_p >= 0.0 and top_p <= 1.0
+
+        if history is None:
+            history = []
+        
+        device = next(self.model.parameters()).device
+        cache_manager = KVCacheManager(
+            num_layers=self.config.n_layer, 
+            batch_size=1, 
+            max_len=self.config.m_len, 
+            num_heads=self.config.n_kvhead,
+            head_dim=self.config.n_dim // self.config.n_head,
+            device=device,
+        )
+        ids = self.tokenizer.encode(build_prompt(query, history))
+        input_ids = torch.tensor([ids], device=device, dtype=torch.long)
+        cpy_history = history.copy()
+        
+        start_cache_pos = len(ids)
+        cur_cache_pos = 0
+        self.model.eval()
+        
+        
+        while len(ids) < self.config.m_len:
+            kv_caches = cache_manager.get_kvcache()
+            logits, cache_increase = self.compute_logits(
+                input_ids,
+                kv_caches=kv_caches, 
+                start_pos=cur_cache_pos
+            )
+            logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
+            probs = torch.softmax(logits, dim=-1)
+            next_token_id = torch.multinomial(probs, num_samples=1)
+            
+            input_ids = next_token_id
+            ids.append(next_token_id.item())
+            cur_cache_pos += cache_increase
+            
+            response = self.tokenizer.decode(ids[start_cache_pos:])
+            yield response, cpy_history + [(query, response)]
+            
+            if next_token_id.item() in self.tokenizer.stop_ids:
+                yield response + "\n", cpy_history + [(query, response)]
+                break
+    
+
+class BatchGenerator(GeneratorCore):
+    def __init__(self, parameter: ModelParameter):
+        super().__init__(parameter)
+    
+    def generate(
+        self, 
+        queries: List[str],
+        histories: List[List[Tuple[str, str]]],
+        temperature: float, 
+        top_k: int, 
+        top_p: float 
+    ) -> List[str]:
+        
+        assert temperature >= 0.0
+        assert top_k >= 0
+        assert top_p >= 0.0 and top_p <= 1.0
+        
+        batch_size = len(queries)
+        if histories is None:
+            histories = [[] for _ in range(batch_size)]
+
+        prompts = [build_prompt(query, history) for query, history in zip(queries, histories)]
+        ids_list = [self.tokenizer.encode(prompt) for prompt in prompts]
+        max_ids_len = max(len(ids) for ids in ids_list)
+        ids_list = pad_sequence(ids_list, max_ids_len, self.tokenizer.pad_id)
+        
+        device = next(self.model.parameters()).device
+        cache_manager = KVCacheManager(
+            num_layers=self.config.n_layer, 
+            batch_size=batch_size, 
+            max_len=self.config.m_len, 
+            num_heads=self.config.n_kvhead,
+            head_dim=self.config.n_dim // self.config.n_head,
+            device=device,
+        )
+        
+        input_tensor = torch.tensor(ids_list, device=device, dtype=torch.long)
+        cache_manager.set_seq_mask(input_tensor, self.tokenizer.pad_id)
+        activate_task_mask = [True] * batch_size
+        
+        start_cache_pos = max_ids_len
+        cur_cache_pos = 0
+        
+        while max_ids_len < self.config.m_len and sum(activate_task_mask) != 0:
+            kv_caches = cache_manager.get_kvcache()
+            attn_mask =cache_manager.get_seq_mask()
+            
+            logits, cache_increase = self.compute_logits(
+                input_tensor,
+                attn_mask=attn_mask,
+                kv_caches=kv_caches, 
+                start_pos=cur_cache_pos
+            )
+            
+            cur_cache_pos += cache_increase
+            logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
+            probs = torch.softmax(logits, dim=-1)
+            next_token_id = torch.multinomial(probs, num_samples=1)
+            
+            active_mask = []
+            c_ids = 0
+            
+            for i in range(batch_size):
+                if activate_task_mask[i]:
+                    token = next_token_id[c_ids, :].item()
+                    ids_list[i].append(token)
+                    c_ids += 1
+                    
+                    is_active = not token in self.tokenizer.stop_ids
+                    activate_task_mask[i] = is_active 
+                    active_mask.append(is_active)
+            
+            active_mask = torch.tensor(active_mask, device=device, dtype=torch.bool)
+            cache_manager.update(active_mask)
+            input_tensor = next_token_id[active_mask, :]
+
+            max_ids_len += 1
+        
+        
+        responses = [str()] * batch_size
+        for i in range(batch_size):
+            responses[i] = self.tokenizer.decode(ids_list[i][start_cache_pos:])
+            histories[i].append((queries[i], responses[i]))
+        
+        return responses
+        
+
+
+class RetrievalGenerator(GeneratorCore):
+    def __init__(self, retriever_parameter: ModelParameter):
+        super().__init__(retriever_parameter)
+    
+    def generate(
+        self,
+        retrieved: List[str],
+        query: str, 
+        history: List[Tuple[str, str]],
+        temperature: float,
+        top_k: int,
+        top_p: float,
+    ) -> str:
+        assert temperature >= 0.0
+        assert top_k >= 0
+        assert top_p >= 0.0 and top_p <= 1.0
+        
+        if history is None:
+            history = []
+            
+        retrieved = "\n".join([f"{idx + 1}. {key}" for idx, key in enumerate(retrieved)]) if retrieved else ""
+        retrieved_query = f"{retrieved}<eos>\n\n根据以上内容回答: {query}" if retrieved else query
+        parameter = ModelParameter(self.model, self.tokenizer, self.config)
+        
+        return ChatGenerator(parameter).generate(
+            retrieved_query, 
+            history,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+        )
+
+class EmbeddingEncoder(EmbeddingEncoderCore):
+    def __init__(self, parameter: ModelParameter):
+        super().__init__(parameter)
+        
+    def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
+        return super().encode(sentence)
+        

khaosz/core/parameter.py

@@ -0,0 +1,238 @@
+import pickle as pkl
+import matplotlib.pyplot as plt
+import safetensors.torch as st
+import torch.nn as nn
+import torch.optim as optim
+
+from dataclasses import dataclass, field
+from typing import Optional, Self, Union
+from pathlib import Path
+
+from khaosz.core.tokenizer import BpeTokenizer
+from khaosz.core.transformer import TransformerConfig, Transformer
+
+
+class BaseModelIO:
+    """Base class for model I/O operations."""
+    
+    def __init__(
+        self, 
+        model: Optional[nn.Module] = None, 
+        tokenizer: Optional[BpeTokenizer] = None,
+        config: Optional[TransformerConfig] = None
+    ):
+        self.model = model
+        self.tokenizer = tokenizer or BpeTokenizer()
+        self.config = config or TransformerConfig()
+    
+    def _get_file_paths(self, directory: Union[str, Path]) -> dict[str, Path]:
+        """Get standardized file paths for model components."""
+        dir_path = Path(directory)
+        return {
+            "model": dir_path / "model.safetensors",
+            "config": dir_path / "config.json", 
+            "tokenizer": dir_path / "tokenizer.json"
+        }
+    
+    def save_components(self, save_dir: Union[str, Path]):
+        """Save core model components."""
+        paths = self._get_file_paths(save_dir)
+        paths["model"].parent.mkdir(parents=True, exist_ok=True)
+        
+        if self.model is not None:
+            st.save_file(self.model.state_dict(), str(paths["model"]))
+        self.config.save(str(paths["config"]))
+        self.tokenizer.save(str(paths["tokenizer"]))
+    
+    def load_components(self, load_dir: Union[str, Path]) -> Self:
+        """Load core model components."""
+        paths = self._get_file_paths(load_dir)
+        
+        self.config.load(str(paths["config"]))
+        self.tokenizer.load(str(paths["tokenizer"]))
+        
+        if paths["model"].exists():
+            state_dict = st.load_file(str(paths["model"]))
+            if self.model is None:
+                self.model = Transformer(self.config)
+            self.model.load_state_dict(state_dict)
+        
+        return self
+    
+    def to(self, *args, **kwargs) -> Self:
+        """Move model to device."""
+        if self.model is not None:
+            self.model.to(*args, **kwargs)
+        return self
+
+
+@dataclass
+class ModelParameter(BaseModelIO):
+    """Container for model parameters with serialization capabilities."""
+    
+    model: Optional[nn.Module] = field(
+        default=None,
+        metadata={"help": "Transformer model."}
+    )
+    tokenizer: BpeTokenizer = field(
+        default_factory=BpeTokenizer,
+        metadata={"help": "Tokenizer for the model."}
+    )
+    config: TransformerConfig = field(
+        default_factory=TransformerConfig,
+        metadata={"help": "Transformer model configuration."}
+    )
+    
+    def save(self, save_dir: Union[str, Path]):
+        """Save model parameters."""
+        self.save_components(save_dir)
+    
+    def load(self, load_dir: Union[str, Path]) -> Self:
+        """Load model parameters."""
+        return self.load_components(load_dir)
+
+
+@dataclass
+class Checkpoint(BaseModelIO):
+    """Extended model parameters with training state."""
+    
+    model: Optional[nn.Module] = field(
+        default=None,
+        metadata={"help": "Transformer model."}
+    )
+    tokenizer: BpeTokenizer = field(
+        default_factory=BpeTokenizer,
+        metadata={"help": "Tokenizer for the model."}
+    )
+    config: TransformerConfig = field(
+        default_factory=TransformerConfig,
+        metadata={"help": "Transformer model configuration."}
+    )
+    loss_list: list[float] = field(
+        default_factory=list,
+        metadata={"help": "List of training losses."}
+    )
+    current_iter: int = field(
+        default=0,
+        metadata={"help": "Current training iteration."}
+    )
+    optimizer: Optional[optim.Optimizer] = field(
+        default=None,
+        metadata={"help": "Optimizer state."}
+    )
+    
+    def __post_init__(self):
+        # Ensure current_iter matches loss list length if not explicitly set
+        if self.current_iter == 0 and self.loss_list:
+            self.current_iter = len(self.loss_list)
+    
+    def _get_training_paths(self, directory: Union[str, Path]) -> dict[str, Path]:
+        """Get file paths for training-specific files."""
+        paths = self._get_file_paths(directory)
+        paths.update({
+            "loss_list": paths["model"].parent / "loss.pkl",
+            "loss_plot": paths["model"].parent / "loss.png",
+            "optimizer": paths["model"].parent / "optimizer.pkl"
+        })
+        return paths
+    
+    def save_training_state(self, save_dir: Union[str, Path]):
+        """Save training-specific state."""
+        paths = self._get_training_paths(save_dir)
+        
+        # Save loss plot
+        self._plot_loss(str(paths["loss_plot"]))
+        
+        # Save loss list
+        with open(str(paths["loss_list"]), "wb") as f:
+            pkl.dump(self.loss_list, f)
+        
+        # Save optimizer state
+        if self.optimizer is not None:
+            with open(str(paths["optimizer"]), "wb") as f:
+                pkl.dump(self.optimizer.state_dict(), f)
+    
+    def load_training_state(self, load_dir: Union[str, Path]) -> Self:
+        """Load training-specific state."""
+        paths = self._get_training_paths(load_dir)
+        
+        # Load loss list
+        if paths["loss_list"].exists():
+            with open(str(paths["loss_list"]), "rb") as f:
+                self.loss_list = pkl.load(f)
+            self.current_iter = len(self.loss_list)
+        
+        # Load optimizer state
+        if paths["optimizer"].exists() and self.optimizer is not None:
+            with open(str(paths["optimizer"]), "rb") as f:
+                optim_state = pkl.load(f)
+            self.optimizer.load_state_dict(optim_state)
+        
+        return self
+    
+    def _plot_loss(self, save_path: str): 
+        """Plot and save loss curve."""
+        if not self.loss_list:
+            return
+            
+        plt.figure(figsize=(10, 6))
+        plt.plot(self.loss_list)
+        plt.title(f"Training Loss - Iteration {self.current_iter}")
+        plt.xlabel("Batch")
+        plt.ylabel("Loss")
+        plt.grid(True)
+        plt.savefig(save_path, dpi=300, bbox_inches="tight")
+        plt.close()
+    
+    def save(self, save_dir: Union[str, Path]):
+        """Save complete checkpoint."""
+        self.save_components(save_dir)
+        self.save_training_state(save_dir)
+    
+    def load(self, load_dir: Union[str, Path]) -> Self:
+        """Load complete checkpoint."""
+        self.load_components(load_dir)
+        self.load_training_state(load_dir)
+        return self
+
+
+class ParameterLoader:
+    """Factory class for loading model parameters or checkpoints."""
+    
+    @staticmethod
+    def load(load_dir: Union[str, Path]) -> Union[ModelParameter, Checkpoint]:
+        """Load either ModelParameter or Checkpoint based on directory contents."""
+        load_dir = Path(load_dir)
+        
+        # Check for training-specific files
+        loss_file = load_dir / "loss.pkl"
+        has_training_data = loss_file.exists()
+        
+        # Create appropriate instance
+        if has_training_data:
+            checkpoint = Checkpoint()
+            checkpoint.load(str(load_dir))
+            return checkpoint
+        else:
+            params = ModelParameter()
+            params.load(str(load_dir))
+            return params
+    
+    @staticmethod
+    def create_checkpoint(
+        model: nn.Module, 
+        tokenizer: BpeTokenizer,
+        config: TransformerConfig,
+        loss_list: Optional[list[float]] = None,
+        optimizer: Optional[optim.Optimizer] = None
+    ) -> Checkpoint:
+        """Convenience method to create a training checkpoint."""
+        return Checkpoint(
+            model=model,
+            tokenizer=tokenizer,
+            config=config,
+            loss_list=loss_list or [],
+            optimizer=optimizer
+        )
+
+

khaosz/core/tokenizer.py

@@ -0,0 +1,111 @@
+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>")

khaosz/core/transformer.py

@@ -0,0 +1,341 @@
+import json
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch import Tensor
+from torch.nn import init
+from dataclasses import asdict, dataclass
+from typing import List, Optional, Self, Tuple
+
+
+def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
+    """ 
+    Repeat k times along the dimension for attention heads.
+    Args:
+        x (Tensor): The input tensor.
+        n_rep (int): The number of repetitions.
+    Returns:
+        Tensor: The repeated tensor.
+    """
+    
+    bs, slen, n_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        x[:, :, :, None, :]
+        .expand(bs, slen, n_heads, n_rep, head_dim)
+        .reshape(bs, slen, n_heads * n_rep, head_dim)
+    )
+
+def get_rotary_emb(
+        dim: int, 
+        max_len: int, 
+        base: float = 10000, 
+        device: torch.device = "cuda",
+    ) -> torch.Tensor:
+    """ 
+    Get the rotary embedding for the given dimension and maximum length.
+    Args:
+        dim (int): The dimension of the input.
+        max_len (int): The maximum length of the input.
+        base (float, optional): The base for the frequency. Defaults to 10000.
+        device (torch.device, optional): The device to use. Defaults to "cuda".
+    Returns:
+        Tensor: The rotary embedding tensor.
+    """
+
+    theta = base ** (-torch.arange(0, dim, 2, device=device).float() / dim)
+    t = torch.arange(0, max_len, device=device).float()
+    freqs = torch.outer(t, theta)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+    
+    return freqs_cis
+
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
+    """
+    Apply rotary embedding to the input tensor.
+    Args:
+        x (Tensor): The input tensor.
+        freqs_cis (Tensor): The rotary embedding tensor.
+    Returns:
+        Tensor: The output tensor.
+    """
+    
+    dtype = x.dtype
+    seq_len = x.size(1)
+
+    x_complex = torch.view_as_complex(x.view(*x.shape[:-1], -1, 2).float())
+    freqs_cis = freqs_cis.reshape(1, seq_len, 1, -1)
+    x_out = torch.view_as_real(x_complex * freqs_cis).flatten(3)
+    
+    return x_out.to(dtype)
+
+def create_attention_mask(
+        seq_mask: Tensor, 
+        start_pos: int = 0,
+        seq_len: int = 0,
+        is_causal: bool = False,
+        device: torch.device = "cuda", 
+        dtype: torch.dtype = torch.float32
+    ) -> Tensor:
+    """
+    Create attention mask for GQA
+    Args:
+        seq_mask (Tensor): A tensor indicating whether each position is valid or not.
+        start_pos (int): The starting position of the sequence.
+        seq_len (int): The length of the sequence.
+        is_causal (bool): Whether the attention is causal or not.
+        device (torch.device): The device to use.
+    Returns:
+        Tensor: The attention mask tensor.
+    """
+    
+    if start_pos != 0 and seq_mask is None:
+        # for single prompt chat
+        seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
+    
+    if seq_mask is None:
+        return None
+    
+    batch_size = seq_mask.size(0)
+    seq_mask = seq_mask[:, :start_pos + seq_len].to(device=device, dtype=torch.bool)
+    # (bsz, start_pos + seq_len)
+    expanded_mask = seq_mask.unsqueeze(1).expand(batch_size, seq_len, start_pos + seq_len)
+    # (bsz, seq_len, start_pos + seq_len)
+    
+    if is_causal:
+        causal_mask = torch.tril(
+            torch.ones((seq_len, start_pos + seq_len), dtype=torch.bool, device=device), 
+            diagonal=start_pos
+        )
+        causal_mask = causal_mask.unsqueeze(0).expand(batch_size, seq_len, start_pos + seq_len)
+        expanded_mask = expanded_mask & causal_mask
+    
+    attention_mask = torch.zeros_like(expanded_mask, dtype=dtype, device=device)
+    attention_mask = attention_mask.masked_fill_(~expanded_mask, -torch.finfo(dtype).max / 2).unsqueeze(1)
+    # (bsz, 1, seq_len, seq_len + start_pos)
+    
+    return attention_mask
+
+
+@dataclass
+class TransformerConfig:
+    # basic config
+    vocab_size: Optional[int] = None
+    n_dim: Optional[int] = None
+    n_head: Optional[int] = None
+    n_layer: Optional[int] = None
+    m_len: Optional[int] = None
+    norm_eps: Optional[float] = None
+    d_ffn: Optional[int] = None
+    
+    # GQA
+    n_kvhead: Optional[int] = None
+    
+    
+    def load(self, config_path: str) -> Self:
+        with open(config_path, 'r') as f:
+            config: dict = json.load(f)
+        for key, value in config.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+    
+        return self
+                    
+    def save(self, config_path: str) -> None:
+        config_dict = asdict(self)
+        config_dict = {k: v for k, v in config_dict.items() if v is not None}
+        with open(config_path, 'w') as f:
+            json.dump(config_dict, f, indent=4)
+            
+
+class Linear(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, bias: bool=False):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
+        self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
+        init.normal_(self.weight, mean=0, std=0.006)
+        
+    def forward(self, x: Tensor) -> Tensor:
+        return F.linear(x, self.weight, self.bias)
+        
+
+class RMSNorm(nn.Module):
+    def __init__(self, n_dim, norm_eps):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(n_dim))
+        self.norm_eps = norm_eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        dtype = x.dtype
+        x = x.float()
+        mean_square = torch.mean(torch.pow(x, 2), dim=-1, keepdim=True)
+        norm = x * torch.rsqrt(mean_square + self.norm_eps)
+        norm = norm.to(dtype)
+        out = norm * self.weight
+        return out
+    
+    
+class MLP(nn.Module):
+    def __init__(self, n_dim: int, d_ffn: int):
+        super().__init__()
+        self.up = Linear(n_dim, d_ffn)
+        self.gate = Linear(n_dim, d_ffn)
+        self.down = Linear(d_ffn, n_dim)
+    
+    def forward(self, x: Tensor) -> Tensor:
+        gated = self.up(x) * F.silu(self.gate(x))
+        out = self.down(gated)
+        return out
+
+
+class GQA(nn.Module):
+    def __init__(
+        self, 
+        n_dim: int, 
+        n_head: int, 
+        n_kvhead: int, 
+    ):
+        super().__init__()
+        assert n_dim % n_head == 0
+        assert n_head % n_kvhead == 0
+        
+        self.head_dim = n_dim // n_head
+        self.n_dim = n_dim
+        self.n_heads = n_head
+        self.n_kvheads = n_kvhead
+        self.n_rep = n_head // n_kvhead
+        
+        self.q_proj = Linear(n_dim, n_head * self.head_dim)
+        self.k_proj = Linear(n_dim, n_kvhead * self.head_dim)
+        self.v_proj = Linear(n_dim, n_kvhead * self.head_dim)
+        self.o_proj = Linear(n_dim, n_dim)
+    
+    def forward(
+        self,
+        x: Tensor, 
+        freqs_cis: Tensor, 
+        mask: Tensor = None,
+        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
+        start_pos: int = 0
+    ) -> Tensor:
+        bsz, seq_len, _ = x.size()
+        # x(bsz, seq_len, n_heads * head_dim) -> (bsz, seq_len, n_heads, head_dim)
+        q = self._split_heads(self.q_proj(x), self.n_heads)
+        k = self._split_heads(self.k_proj(x), self.n_kvheads)
+        v = self._split_heads(self.v_proj(x), self.n_kvheads)
+        q, k = apply_rotary_emb(q, freqs_cis), apply_rotary_emb(k, freqs_cis)
+        
+        if kv_cache is not None:
+            k_cache, v_cache = kv_cache
+
+            # copy to cache
+            k_cache[:bsz, start_pos:start_pos + seq_len] = k
+            v_cache[:bsz, start_pos:start_pos + seq_len] = v
+            
+            # get cache
+            k = k_cache[:bsz, :start_pos + seq_len]
+            v = v_cache[:bsz, :start_pos + seq_len]
+        
+        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
+        
+        # (bsz, seq_len, n_heads, head_dim) -> (bsz, n_heads, seq_len, head_dim)
+        q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
+        sdqa_out = F.scaled_dot_product_attention(q, k, v, mask, is_causal=(mask == None)).permute(0, 2, 1, 3)
+        out = self.o_proj(sdqa_out.contiguous().view(bsz, seq_len, -1))
+
+        return out
+    
+    def _split_heads(self, x: Tensor, n_heads) -> Tensor:
+        batch_size, seq_len, _ = x.shape
+        x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
+        return x
+    
+
+class DecoderBlock(nn.Module):
+    def __init__(self, n_dim, n_head, d_ffn, n_kvhead, norm_eps):
+        super().__init__()
+        self.attention = GQA(n_dim, n_head, n_kvhead)
+        self.norm_attn = RMSNorm(n_dim, norm_eps)
+        self.ffn = MLP(n_dim, d_ffn)
+        self.norm_ffn = RMSNorm(n_dim, norm_eps)
+
+    def forward(
+        self,
+        x: Tensor,
+        freqs_cis: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
+        start_pos: int = 0
+    ) -> Tensor:
+        # attention
+        attn_output = self.attention(
+            self.norm_attn(x), 
+            freqs_cis, 
+            attention_mask, 
+            kv_cache, 
+            start_pos
+        )
+        x = attn_output + x
+        
+        # feed forward
+        x = self.ffn(self.norm_ffn(x)) + x
+        
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        self.embedding = nn.Parameter(torch.empty(config.vocab_size, config.n_dim))
+        self.layers = nn.ModuleList([
+            DecoderBlock(
+                config.n_dim, 
+                config.n_head, 
+                config.d_ffn, 
+                config.n_kvhead, 
+                config.norm_eps
+            )
+            for _ in range(config.n_layer)
+        ])
+        self.norm = RMSNorm(config.n_dim, config.norm_eps)
+        self.freq_cis = get_rotary_emb(config.n_dim // config.n_head, config.m_len)
+        init.normal_(self.embedding, mean=0, std=0.02)
+    
+    def forward(
+        self, 
+        input_ids: Tensor, 
+        seq_mask: Optional[Tensor]=None,
+        persistent_key_values: Optional[List[Tuple[Tensor, Tensor]]]=None,
+        start_pos: int = 0
+    ) -> Tensor:
+        assert input_ids.ndim == 2
+        seq_len = input_ids.size(-1)
+        x = F.embedding(input_ids, self.embedding)
+        
+        self.freq_cis = self.freq_cis.to(x.device)
+        freqs_cis = self.freq_cis[start_pos:start_pos+seq_len]
+        has_kvcache = persistent_key_values is not None
+        
+        attn_mask = create_attention_mask(
+            seq_mask, 
+            start_pos=start_pos,
+            seq_len=seq_len,
+            is_causal=has_kvcache,
+            device=x.device,
+            dtype=x.dtype
+        )
+        
+        for i, layer in enumerate(self.layers):
+            kv_cache = persistent_key_values[i] if persistent_key_values else None
+            x = layer(x, freqs_cis, attn_mask, kv_cache, start_pos)
+        
+        hidden_states = self.norm(x)        
+        logits = F.linear(hidden_states,  self.embedding)
+        
+        return {
+            "logits": logits,
+            "hidden_states": hidden_states
+        }
+        

khaosz/model.py

@@ -0,0 +1,112 @@
+from torch import Tensor
+from typing import List, Tuple, Generator, Union
+
+from khaosz.core.generator import (
+    TextGenerator,
+    ChatGenerator, 
+    StreamGenerator, 
+    BatchGenerator, 
+    RetrievalGenerator, 
+    EmbeddingEncoder
+)
+from khaosz.core.parameter import ParameterLoader
+
+
+class Khaosz:
+    def __init__(self, model_dir: str):
+        self.parameter = ParameterLoader.load(model_dir)
+    
+    def to(self, *args, **kwargs):
+        self.parameter.to(*args, **kwargs)
+        return self
+
+    def generate(
+            self, 
+            query: str, 
+            history: List[Tuple[str, str]]=None,
+            temperature: float=0.8,
+            top_k: int=50,
+            top_p: float=0.95,
+        ) -> str:
+            generator = ChatGenerator(self.parameter)
+            return generator.generate(
+                query, 
+                history=history,
+                temperature=temperature, 
+                top_k=top_k, 
+                top_p=top_p,
+            )
+    
+    def batch_generate(
+            self, 
+            queries: List[str],
+            histories: List[Tuple[str, str]]=None,
+            temperature: float=0.8,
+            top_k: int=50,
+            top_p: float=0.95,
+        ) -> List[str]:
+            generator = BatchGenerator(self.parameter)
+            return generator.generate(
+                queries, 
+                histories=histories,
+                temperature=temperature, 
+                top_k=top_k, 
+                top_p=top_p,
+            )
+        
+    
+    def stream_generate(
+            self, 
+            query: str, 
+            history: List[Tuple[str, str]]=None,
+            temperature: float=0.8,
+            top_k: int=50,
+            top_p: float=0.95,
+        ) -> Generator[Tuple[str, List[Tuple[str, str]]], None, None]:
+            stream_generator = StreamGenerator(self.parameter)
+            return stream_generator.generate(
+                    query, 
+                    history=history,
+                    temperature=temperature, 
+                    top_k=top_k, 
+                    top_p=top_p,
+                )
+    
+    def retrieve_generate(
+            self,
+            retrieved,
+            query: str, 
+            history: List[Tuple[str, str]] = None,
+            temperature: float=0.8,
+            top_k: int=50,
+            top_p: float=0.95,
+        ) -> str:
+            generator = RetrievalGenerator(self.parameter)
+            return generator.generate(
+                retrieved,
+                query,
+                history=history,
+                temperature=temperature, 
+                top_k=top_k, 
+                top_p=top_p,
+            )
+    
+    def text_generate(            
+            self, 
+            query: str, 
+            temperature: float=0.8,
+            top_k: int=50,
+            top_p: float=0.95,
+        ) -> str:
+        generator = TextGenerator(self.parameter)
+        
+        return generator.generate(
+                query,
+                temperature=temperature, 
+                top_k=top_k, 
+                top_p=top_p,
+            )
+    
+    def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
+        encoder = EmbeddingEncoder(self.parameter)
+        return encoder.encode(sentence)

khaosz/trainer/__init__.py

@@ -0,0 +1,11 @@
+from khaosz.trainer.dataset import DatasetLoader
+from khaosz.trainer.trainer import Trainer
+from khaosz.trainer.strategy import TrainConfig, CosineScheduleConfig, SgdrScheduleConfig
+
+__all__ = [
+    "DatasetLoader",
+    "Trainer",
+    "TrainConfig",
+    "CosineScheduleConfig",
+    "SgdrScheduleConfig",
+]

khaosz/trainer/dataset.py

@@ -0,0 +1,210 @@
+import torch
+import bisect
+import pickle as pkl
+from abc import ABC, abstractmethod
+from torch import Tensor
+from torch.utils.data import Dataset
+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
+
+
+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)
+        
+        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx - 1)
+        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx - 1)
+
+        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, device: str):
+        super().__init__()
+        self.segments: MutiSeg = {}
+        self.chunk_size = chunk_size
+        self.total_samples = 0
+        self.device = device
+
+    def save(self, save_path: str):      
+        first_item = self.segments[keys[0]]
+        segment_size = len(first_item)
+        keys = list(self.segments.keys())
+        
+        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):
+        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 , device='cuda'):
+        super().__init__(chunk_size, device)
+        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):
+        begin_idx = index * self.chunk_size 
+        end_idx = min(begin_idx + self.chunk_size, self.total_samples - 1)
+        
+        x = self._fetch_data(begin_idx, end_idx).to(device=self.device, dtype=torch.long)
+        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(device=self.device, dtype=torch.long)
+        
+        return x, y
+    
+    
+class SftDataset(BaseDataset):
+    def __init__(self, chunk_size, device='cuda'):
+        super().__init__(chunk_size, device)
+        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):
+        begin_idx = index * self.chunk_size 
+        end_idx = min(begin_idx + self.chunk_size, self.total_samples - 1)
+        
+        x = self._fetch_data(begin_idx, end_idx, "sequence").to(device=self.device, dtype=torch.long)
+        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence").to(device=self.device, dtype=torch.long)
+        loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "mask").to(device=self.device, dtype=torch.bool)
+        
+        return x, y, loss_mask
+
+
+class DpoDataset(BaseDataset):
+    def __init__(self, chunk_size: int, device="cuda"):
+        super().__init__(chunk_size, device)
+        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):
+        start_idx = index * self.chunk_size
+        end_idx = min(start_idx + self.chunk_size, self.total_samples - 1)
+        
+        chosen = self._fetch_data(start_idx, end_idx, "chosen").to(device=self.device, dtype=torch.long)
+        rejected = self._fetch_data(start_idx, end_idx, "rejected").to(device=self.device, dtype=torch.long)
+        chosen_mask = self._fetch_data(start_idx, end_idx, "chosen_mask").to(device=self.device, dtype=torch.bool)
+        rejected_mask = self._fetch_data(start_idx, end_idx, "rejected_mask").to(device=self.device, dtype=torch.bool)
+
+        return chosen, rejected, chosen_mask, rejected_mask
+
+
+class PpoDataset(BaseDataset):
+    def __init__(self, chunk_size: int, device="cuda"):
+        super().__init__(chunk_size, device)
+        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 = index * self.chunk_size
+        end_idx = min(begin_idx + self.chunk_size, self.total_samples - 1)
+        
+
+        input_ids =  self._fetch_data(begin_idx, end_idx, "input_ids").to(self.device),
+        actions = self._fetch_data(begin_idx, end_idx, "actions").to(self.device),
+        logprobs = self._fetch_data(begin_idx, end_idx, "logprobs").to(self.device),
+        rewards =  self._fetch_data(begin_idx, end_idx, "rewards").to(self.device)
+        
+        return input_ids, actions, logprobs, rewards
+    
+
+class DatasetLoader:
+    @staticmethod       
+    def load(
+        train_type: Literal["seq", "sft", "dpo"],
+        load_path: Union[str, List[str]],
+        max_len: int, 
+        device: str
+        ) -> BaseDataset:
+        
+        dataset_router: Dict[str, Callable[[int, torch.device], BaseDataset]] = {
+            "seq": lambda m_len, device: SeqDataset(m_len, device=device),
+            "sft": lambda m_len, device: SftDataset(m_len, device=device),
+            "dpo": lambda m_len, device: DpoDataset(m_len, device=device),
+        }
+        dataset = dataset_router[train_type](max_len, device)
+        dataset.load(load_path)
+        
+        return dataset

khaosz/trainer/mask.py

@@ -0,0 +1,55 @@
+
+import torch
+from abc import abstractmethod
+from torch import Tensor
+
+
+
+class MaskBuilder:
+    def __init__(
+        self,
+        bos_token_id: int,
+        eos_token_id: int,
+        user_token_id: int,
+        system_token_id: int,
+        
+    ):
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.user_token_id = user_token_id
+        self.system_token_id = system_token_id
+    
+    @abstractmethod
+    def build(input_ids: Tensor) -> Tensor:
+        raise NotImplementedError
+
+
+
+class LossMaskBuilder(MaskBuilder):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+    
+    def build(self, input_ids: Tensor) -> Tensor:
+        token_markers = torch.zeros_like(input_ids, dtype=torch.int8)
+        
+        is_user_token = input_ids.eq(self.user_token_id)
+        is_system_token = input_ids.eq(self.system_token_id)
+        
+        token_markers[is_user_token] = 1
+        token_markers[is_system_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
+        
+        
+        
+
+class AttentionMaskBuilder:
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+    
+    def build(input_ids: Tensor):
+        bsz = input_ids.size(0)

khaosz/trainer/strategy.py

@@ -0,0 +1,388 @@
+import copy
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch import Tensor
+from torch.optim import Optimizer
+from torch.utils.data import Dataset
+from typing import Any, Literal, Tuple, Callable, Dict
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass, field
+
+
+def get_logprobs(model:nn.Module, input_ids: Tensor, mask: Tensor, pad_token_id):
+    input_mask =  input_ids.ne(pad_token_id)
+    logits = model(input_ids, input_mask)["logits"]
+    log_probs = torch.log_softmax(logits, dim=-1)
+    
+    shifted_log_probs = log_probs[:, :-1, :] 
+    shifted_input_ids = input_ids[:, 1:]
+    shifted_response_mask = mask[:, 1:]
+    
+    token_logprobs = torch.gather(
+        shifted_log_probs, 
+        dim=-1, 
+        index=shifted_input_ids.unsqueeze(-1)
+    ).squeeze(-1)
+    
+    prompt_mask = input_mask[:, 1:]
+    valid_mask = (prompt_mask & shifted_response_mask).float()
+    
+    return (token_logprobs * valid_mask).sum(dim=-1)
+
+
+class MaskBuilder:
+    def __init__(
+        self,
+        bos_token_id: int,
+        eos_token_id: int,
+        user_token_id: int,
+        system_token_id: int,
+        
+    ):
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.user_token_id = user_token_id
+        self.system_token_id = system_token_id
+    
+    @abstractmethod
+    def build(input_ids: Tensor) -> Tensor:
+        raise NotImplementedError
+
+
+
+class LossMaskBuilder(MaskBuilder):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+    
+    def build(self, input_ids: Tensor) -> Tensor:
+        token_markers = torch.zeros_like(input_ids, dtype=torch.int8)
+        
+        is_user_token = input_ids.eq(self.user_token_id)
+        is_system_token = input_ids.eq(self.system_token_id)
+        
+        token_markers[is_user_token] = 1
+        token_markers[is_system_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 AttentionMaskBuilder(MaskBuilder):
+    def __init__(self, multi_turn=False, **kwargs):
+        super().__init__(**kwargs)
+        self.multi_turn = multi_turn
+    
+    def build(self, input_ids: Tensor):
+        bsz = input_ids.size(0)
+        
+    
+    def _build_batch(self, input_ids: Tensor):
+        is_user_token = input_ids.eq(self.user_token_id)
+        
+        token_markers = torch.zeros_like(input_ids, dtype=torch.int8)
+        token_markers[is_user_token] = 1
+        cumulative_markers = torch.cumsum(token_markers, dim=-1)
+            
+            
+        
+
+
+class BaseStrategy(ABC):
+    def __init__(self, model: nn.Module):
+        self.model = model
+    
+    @abstractmethod
+    def compute_loss(self, batch: Tuple[Tensor, ...]) -> Tensor:
+        raise NotImplementedError
+    
+    def __call__(self, batch: Tuple[Tensor, ...]) -> Tensor:
+        return self.compute_loss(batch)
+
+
+class SeqStrategy(BaseStrategy):
+    def __init__(self, model):
+        super().__init__(model)
+    
+    def compute_loss(self, batch: Tuple[Tensor, ...]) -> Tensor:
+        x, y = batch
+        B, L = x.size()
+        logits: Tensor = self.model(x)["logits"]
+        
+        loss = F.cross_entropy(
+            logits.view(B * L, -1), y.flatten()
+        )
+        return loss
+    
+
+class SftStrategy(BaseStrategy):
+    def __init__(self, model):
+        super().__init__(model)
+    
+    def compute_loss(self, batch: Tuple[Tensor, ...]) -> Tensor:
+        x, y, loss_mask = batch
+        B, L = x.size()
+        ignore_idx = -1
+        
+        logits: Tensor = self.model(x)["logits"]
+        masked_y = y.masked_fill(loss_mask == 0, ignore_idx)
+        
+        loss = F.cross_entropy(
+            logits.view(B * L, -1),
+            masked_y.flatten(), 
+            ignore_index=ignore_idx
+        )
+
+        return loss
+
+class DpoStrategy(BaseStrategy):
+    def __init__(self, model, pad_token_id, beta):
+        super().__init__(model)
+        ref_model = copy.deepcopy(self.model)
+        ref_model.requires_grad_(False)
+        ref_model.eval()
+        
+        self.ref_model = ref_model
+        self.pad_token_id = pad_token_id
+        self.beta = beta
+        
+    def compute_loss(self, batch: Tuple[Tensor, ...]) -> Tensor:
+        good_ids, bad_ids, good_mask, bad_mask = batch
+        
+        log_pi_good = get_logprobs(self.model, good_ids, good_mask, self.pad_token_id)
+        log_pi_bad = get_logprobs(self.model, bad_ids, bad_mask, self.pad_token_id)
+        
+        with torch.no_grad():
+            log_ref_good = get_logprobs(self.ref_model, good_ids, good_mask, self.pad_token_id)
+            log_ref_bad = get_logprobs(self.ref_model, bad_ids, bad_mask, self.pad_token_id)
+        
+        pi_log_ratio = log_pi_good - log_pi_bad
+        ref_log_ratio = log_ref_good - log_ref_bad
+
+        ratio_diff = pi_log_ratio - ref_log_ratio
+        
+        dpo_loss = -F.logsigmoid(self.beta * ratio_diff).mean()
+        return dpo_loss
+
+
+class PpoStrategy(BaseStrategy):
+    def __init__(self, model, pad_token_id, epsilon):
+        super().__init__(model)
+        ref_model = copy.deepcopy(self.model)
+        ref_model.requires_grad_(False)
+        ref_model.eval()
+        
+        self.ref_model = ref_model
+        self.pad_token_id = pad_token_id
+        self.epsilon = epsilon
+        
+    def ppo_clip_loss_masked(
+        self,
+        log_probs: Tensor, 
+        old_log_probs: Tensor, 
+        advantages: Tensor, 
+        values: Tensor, 
+        returns: Tensor,
+        mask: Tensor, 
+        clip_eps: float=0.2, 
+    ):
+        ratio = torch.exp(log_probs - old_log_probs)
+        surr1 = ratio * advantages
+        surr2 = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) * advantages
+        policy_loss = -torch.min(surr1, surr2).masked_select(mask).mean()
+
+        value_loss = F.mse_loss(values.masked_select(mask),
+                                returns.masked_select(mask))
+
+        entropy = -(log_probs.exp() * log_probs).masked_select(mask).mean()
+        entropy_loss = -entropy
+        return policy_loss, value_loss, entropy_loss
+
+
+
+class StrategyFactory:
+    
+    def load(model, train_type, pad_token_id, dpo_beta):
+        train_strategy: Dict[str, Callable[[], BaseStrategy]] = {
+            "seq": lambda: SeqStrategy(model),
+            "sft": lambda: SftStrategy(model),
+            "dpo": lambda: DpoStrategy(model, pad_token_id, dpo_beta)
+        }
+        strategy = train_strategy[train_type]()
+        return strategy
+    
+
+@dataclass
+class TrainConfig:
+    train_type: str = field(
+        default_factory=["seq", "sft", "dpo"],
+        metadata={"help": "Type of training."}
+    )
+    dataset: Dataset = field(
+        default=None,
+        metadata={"help": "Dataset for training."}
+    )
+    optimizer: Optimizer = field(
+        default=None,
+        metadata={"help": "Optimizer for training."}
+    )
+    ckpt_dir: str = field(
+        default="./checkpoint",
+        metadata={"help": "Checkpoint directory."}
+    )
+    n_epoch: int = field(
+        default=1,
+        metadata={"help": "Number of epochs for training."}
+    )
+    batch_size: int = field(
+        default=4,
+        metadata={"help": "Batch size for training."}
+    )
+    n_iter_ckpt: int = field(
+        default=5000,
+        metadata={"help": "Number of iterations between checkpoints."}
+    )
+    n_iter_step: int = field(
+        default=1,
+        metadata={"help": "Number of iterations between steps."}
+    )
+    max_grad_norm: float = field(
+        default=1.0,
+        metadata={"help": "Maximum gradient norm."}
+    )
+    random_seed: int = field(
+        default=3407,
+        metadata={"help": "Random seed."}
+    )
+    dpo_beta: float = field(
+        default=0.1,
+        metadata={"help": "DPO beta."}
+    )
+
+    def get_kwargs(self)-> Dict[str, Any]:
+        config_dict = asdict(self)
+        return {k: v for k, v in config_dict.items() if v is not None}
+    
+
+@dataclass
+class ScheduleConfig:
+    schedule_type: str = field(
+        default_factory=["cosine", "sgdr"],
+        metadata={"help": "Type of learning rate schedule."}
+    )
+    warning_step: int = field(
+        default=1000,
+        metadata= {"help": "Warning up step."}
+    )
+    @abstractmethod
+    def get_kwargs(self)-> Dict[str, Any]:
+        raise NotImplementedError
+
+
+@dataclass   
+class CosineScheduleConfig(ScheduleConfig):
+    total_iters: int = field(
+        default=None,
+        metadata={"help": "Total iterations for cosine schedule."}
+    )
+    min_rate: float = field(
+        default=0.05,
+        metadata={"help": "Minimum rate for cosine schedule."}
+    )
+    schedule_type: Literal["cosine"] = "cosine"
+    
+    def get_kwargs(self) -> Dict[str, Any]:
+        return {
+            "schedule_type": self.schedule_type,
+            "warning_step": self.warning_step,
+            "lr_decay_iters": self.total_iters - self.warning_step,
+            "min_rate": self.min_rate
+        }
+
+@dataclass
+class SgdrScheduleConfig(ScheduleConfig):
+    cycle_length: int = field(
+        default=1000,
+        metadata={"help": "Cycle length for sgdr schedule."}
+    )
+    min_rate: float = field(
+        default=0.05,
+        metadata={"help": "Minimum rate for sgdr schedule."}
+    )
+    T_mult: int = field(
+        default=2,
+        metadata={"help": "T_mult for sgdr schedule."}
+    )
+    schedule_type: Literal["sgdr"] = "sgdr"
+
+    def get_kwargs(self) -> Dict[str, Any]:
+        return {
+            "schedule_type": self.schedule_type,
+            "warning_step": self.warning_step,
+            "cycle_length": self.cycle_length,
+            "min_rate": self.min_rate,
+            "T_mult": self.T_mult
+        }
+
+
+class SchedulerFactory:
+
+    @staticmethod
+    def get_sgdr_schedule(
+        warning_step: int, 
+        cycle_length: int, 
+        min_rate: float = 0.1, 
+        T_mult: int = 2
+    ) -> Callable[[int], float]:
+
+        def sgdr_schedule(now_iter: int) -> float:
+            if now_iter < warning_step:
+                return max(min_rate, now_iter / warning_step)
+                
+            adjusted_iter = now_iter - warning_step
+            total_cycles, current_cycle = 0, 0
+            while adjusted_iter >= cycle_length * (T_mult ** total_cycles):
+                current_cycle += 1
+                total_cycles += 1
+            
+            cycle_start = sum(cycle_length * (T_mult ** i) for i in range(current_cycle))
+            cycle_pos = adjusted_iter - cycle_start
+            
+            cycle_length_current = cycle_length * (T_mult ** current_cycle)
+            return max(min_rate, 0.5 * (1 + math.cos(math.pi * cycle_pos / cycle_length_current)))
+        
+        return sgdr_schedule
+
+    @staticmethod
+    def get_cosine_warmup_schedule(
+        warning_step: int, 
+        lr_decay_iters: int, 
+        min_rate: float = 0.1
+    ) -> Callable[[int], float]:
+
+        def cosine_warmup_schedule(now_iter: int) -> float:
+            if now_iter <= warning_step:
+                return max(min_rate, now_iter / warning_step)
+            else:
+                rate = (now_iter - warning_step) / (lr_decay_iters - warning_step)
+                return max(min_rate, 0.5 * (1.0 + math.cos(math.pi * rate)))
+        
+        return cosine_warmup_schedule
+    
+    @staticmethod
+    def load_schedule_fn(**kwargs):
+        strategy = kwargs.pop("schedule_type")
+        if strategy == "cosine":
+            return SchedulerFactory.get_cosine_warmup_schedule(**kwargs)
+        elif strategy == "sgdr":
+            return SchedulerFactory.get_sgdr_schedule(**kwargs)
+        else:
+            raise ValueError(f"Invalid schedule type: {strategy}")
+        

khaosz/trainer/trainer.py

@@ -0,0 +1,167 @@
+import os
+import torch
+import logging
+
+from typing import Tuple
+from torch.nn.utils import clip_grad_norm_
+from torch.optim.lr_scheduler import LambdaLR
+from torch.utils.data import DataLoader, RandomSampler
+from tqdm import tqdm
+
+from khaosz.core import ModelParameter, Checkpoint
+from khaosz.trainer.strategy import SchedulerFactory, StrategyFactory, TrainConfig, ScheduleConfig
+
+
+class Trainer:
+    def __init__(
+        self,
+        parameter: ModelParameter,
+        log_path: str="./train_log.log"
+    ):
+        logger = logging.getLogger()
+        logger.setLevel(level = logging.INFO)
+        handler = logging.FileHandler(log_path)
+        handler.setLevel(logging.INFO)
+        handler.setFormatter(logging.Formatter('%(asctime)s: %(message)s'))
+        logger.addHandler(handler)
+        logger.info("initializing trainer ...")
+        
+        self.logger = logger
+        self.model = parameter.model
+        self.tokenizer = parameter.tokenizer
+        self.config = parameter.config
+        
+    def save_checkpoint(
+        self, 
+        loss_list: list, 
+        ckpt_dir: str, 
+        current_iter: int, 
+        last_ckpt_iter: int
+    ):
+        save_path = os.path.join(ckpt_dir, f"iter_{current_iter}")
+        Checkpoint(
+            self.model, 
+            self.tokenizer, 
+            self.config, 
+            loss_list, 
+            current_iter
+        ).save(save_path)
+        
+        diff_iter = current_iter - last_ckpt_iter
+        avg_loss = sum(loss_list[last_ckpt_iter:current_iter]) / diff_iter
+        self.logger.info(f"iter: {current_iter} loss: {avg_loss}")  
+        
+        return current_iter
+    
+    def load_checkpoint(self, train_checkpoint: Checkpoint) -> Tuple[list, int]:
+        self.model = train_checkpoint.model
+        self.tokenizer = train_checkpoint.tokenizer
+        self.config = train_checkpoint.config
+        loss_list = train_checkpoint.loss_list
+        last_ckpt_iter = train_checkpoint.current_iter
+        
+        return loss_list, last_ckpt_iter
+
+    def train(
+        self,
+        train_config: TrainConfig,
+        schedule_config: ScheduleConfig,
+        train_checkpoint: Checkpoint = None
+    ):
+        assert schedule_config.schedule_type in ["cosine", "sgdr"]
+        assert train_config.train_type in ["seq", "sft", "dpo"]
+        
+        if train_checkpoint:
+            loss_list, last_ckpt_iter = self.load_checkpoint(train_checkpoint)
+            current_iter = train_checkpoint.current_iter + 1
+            self.logger.info(f"Resuming training from checkpoint: iter {current_iter}")
+        else:
+            current_iter = 0
+            last_ckpt_iter = 0
+            loss_list = []
+            
+        lambda_scheduler_fn  = SchedulerFactory.load_schedule_fn(
+            **schedule_config.get_kwargs()
+        )
+        
+        strategy = StrategyFactory.load(
+            self.model, 
+            train_config.train_type, 
+            self.tokenizer.pad_id, 
+            train_config.dpo_beta
+        )
+        
+        scheduler = LambdaLR(
+            train_config.optimizer,
+            lambda_scheduler_fn,
+            last_epoch=current_iter - 1 if train_checkpoint else -1
+        )
+        
+        seed = train_config.random_seed
+        generator = torch.Generator().manual_seed(seed)
+        sampler = RandomSampler(train_config.dataset, generator=generator)
+        remaining_epochs = train_config.n_epoch - current_iter // (len(train_config.dataset) // train_config.batch_size)
+        
+        self.logger.info(f"Starting {train_config.train_type.upper()} training for {train_config.n_epoch} epochs")
+        self.logger.info(f"Checkpoint interval: {train_config.n_iter_ckpt} iterations")
+        
+        for epoch in range(remaining_epochs):
+            self.model.train()
+            dataloader = DataLoader(
+                train_config.dataset, 
+                batch_size=train_config.batch_size, 
+                sampler=sampler
+            )
+            progress_bar = tqdm(
+                dataloader, 
+                desc=f"Epoch {epoch+1}/{train_config.n_epoch}", 
+                dynamic_ncols=True
+            )
+            for batch in progress_bar:
+                #forward
+                loss = strategy(batch)
+                loss_list.append(loss.item())
+                #backward
+                loss.backward()
+                #step
+                if current_iter % train_config.n_iter_step == 0:
+                    clip_grad_norm_(
+                        self.model.parameters(),
+                        train_config.max_grad_norm
+                    )
+                    train_config.optimizer.step()
+                    train_config.optimizer.zero_grad()
+                    
+                current_iter += 1
+                scheduler.step()
+                progress_bar.set_postfix({
+                    "loss": f"{loss.item():.4f}",
+                    "lr": f"{train_config.optimizer.param_groups[0]['lr']:.2e}"
+                })
+                #save checkpotint
+                if current_iter - last_ckpt_iter >= train_config.n_iter_ckpt:
+                    last_ckpt_iter = self.save_checkpoint(
+                        loss_list, 
+                        train_config.ckpt_dir, 
+                        current_iter, 
+                        last_ckpt_iter
+                    )
+
+        if current_iter != last_ckpt_iter:
+            last_ckpt_iter = self.save_checkpoint(
+                loss_list, 
+                train_config.ckpt_dir, 
+                current_iter, 
+                last_ckpt_iter
+            )
+
+        self.logger.info("Training completed")
+        
+        return Checkpoint(
+            self.model,
+            self.tokenizer,
+            self.config,
+            loss_list,
+            current_iter,
+            train_config.optimizer
+        )

khaosz/utils/retriever.py

@@ -0,0 +1,88 @@
+import torch
+import sqlite3
+import numpy as np
+from torch import Tensor
+from typing import Dict, List, Tuple
+
+
+class Retriever:
+    def __init__(self, db_path=None):
+        self.data: Dict[str, Tensor] = {}
+        self.embedding_cache: Tensor = None
+        self.is_caculated: bool = False
+        
+        if db_path is not None:
+            self.load(db_path)
+                
+    def retrieve(self, query: Tensor, top_k: int) -> List[Tuple[str, float]]:
+        if not self.data:
+            return []
+        
+        query = query.flatten().unsqueeze(1)                            # [dim, 1]
+        norm_embeddings = self._embeddings.to(
+            device=query.device,
+            dtype=query.dtype
+        )                                                               # [n_vectors, dim]
+        sim_scores = torch.matmul(norm_embeddings, query).squeeze()     # [n_vectors]
+        
+        top_k = min(top_k, len(self.data))
+        indices = sim_scores.topk(top_k).indices
+        keys = list(self.data.keys())
+        
+        return [(keys[i], sim_scores[i].item()) for i in indices]
+    
+    def add_vector(self, key: str, vector_data: Tensor):
+        self.is_caculated = False
+        self.data[key] = vector_data.flatten().float().cpu()
+        
+    def delete_vector(self, key: str):
+        self.is_caculated = False
+        self.data.pop(key, None)
+    
+    def save(self, db_path):
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+        self._init_db(cursor)
+        cursor.execute('DELETE FROM vectors')
+        
+        for item, vec in self.data.items():
+            vec_bytes = vec.numpy().tobytes()
+            cursor.execute('INSERT OR REPLACE INTO vectors (key, vector) VALUES (?, ?)', 
+                           (item, vec_bytes))
+        
+        conn.commit()
+        conn.close()
+
+    def load(self, db_path):
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+        self._init_db(cursor)
+        cursor.execute('SELECT key, vector FROM vectors')
+        rows = cursor.fetchall()
+        self.data = {}
+        
+        for row in rows:
+            key, vec_bytes = row
+            vec_numpy = np.frombuffer(vec_bytes, dtype=np.float32).copy()
+            vec = torch.from_numpy(vec_numpy)
+            self.data[key] = vec
+        
+        conn.close()
+        
+    def _init_db(self,cursor: sqlite3.Cursor):
+        # Create table if not exists (in case loading from a new database)
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS vectors (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                key TEXT UNIQUE NOT NULL,
+                vector BLOB NOT NULL
+            )''')
+    
+    @property
+    def _embeddings(self) -> Tensor:
+        if not self.is_caculated:
+            embeddings = torch.stack(list(self.data.values())) 
+            norm_embeddings = embeddings / torch.norm(embeddings, dim=-1, keepdim=True)
+            self.embedding_cache = norm_embeddings
+        
+        return self.embedding_cache

khaosz/utils/splitter.py

@@ -0,0 +1,127 @@
+import re
+import torch
+import torch.nn.functional as F
+
+from abc import ABC, abstractmethod
+from torch import Tensor
+from typing import List, Callable, Optional
+
+
+class BaseTextSplitter(ABC):
+    def __init__(
+        self,
+        max_len: int = 512,
+        chunk_overlap: int = 0,
+    ):
+        if max_len <= 0:
+            raise ValueError("max_len must be > 0")
+        if chunk_overlap < 0:
+            raise ValueError("chunk_overlap must be >= 0")
+
+        self.max_len = max_len
+        self.chunk_overlap = chunk_overlap
+
+    @abstractmethod
+    def split(self, text: str, **kwargs) -> List[str]:
+        raise NotImplementedError
+
+    def preprocess(self, text: str) -> str:
+        return text.strip()
+
+    def postprocess(self, chunks: List[str]) -> List[str]:
+        return [chunk.strip() for chunk in chunks if chunk.strip()]
+
+
+class PriorityTextSplitter(BaseTextSplitter):
+    def __init__(
+        self,
+        separators: List[str],
+        max_len: int = 512,
+        chunk_overlap: int = 0,
+    ):
+        super().__init__(max_len=max_len, chunk_overlap=chunk_overlap)
+        if not separators:
+            raise ValueError("separators must be a non-empty list")
+        self.separators = separators
+
+    def split(self, text: str) -> List[str]:
+        text = self.preprocess(text)
+        for sep in self.separators:
+            parts = text.split(sep)
+
+            valid_parts = [p.strip() for p in parts if p.strip()]
+            if len(valid_parts) > 1:
+                return self.postprocess(valid_parts)
+        return [text]
+
+
+class SemanticTextSplitter(BaseTextSplitter):
+
+    DEFAULT_PATTERN = r'(?<=[。！？!?])(?=(?:[^"\'‘’“”]*["\'‘’“”][^"\'‘’“”]*["\'‘’“”])*[^"\'‘’“”]*$)'
+
+    def __init__(
+        self,
+        embedding_func: Callable[[List[str]], List[Tensor]],
+        pattern: Optional[str] = None,
+        max_len: int = 512,
+        chunk_overlap: int = 0,
+    ):
+        super().__init__(max_len=max_len, chunk_overlap=chunk_overlap)
+        if not callable(embedding_func):
+            raise TypeError("embedding_func must be callable")
+        self.embedding_func = embedding_func
+        self.pattern = pattern or SemanticTextSplitter.DEFAULT_PATTERN
+
+    def split(
+        self,
+        text: str,
+        threshold: float = 0.5,
+        window_size: int = 1,
+    ) -> List[str]:
+        text = self.preprocess(text)
+        sentences = [s.strip() for s in re.split(self.pattern, text) if s.strip()]
+
+        if len(sentences) <= 1:
+            return self.postprocess(sentences)
+
+        try:
+            sentence_embs = self.embedding_func(sentences)
+        except Exception as e:
+            raise RuntimeError(f"Embedding generation failed: {e}")
+
+        if len(sentence_embs) != len(sentences):
+            raise ValueError("Embedding function must return one vector per sentence")
+
+        chunks = []
+        emb_tensor = torch.stack(sentence_embs)  # shape: [N, D]
+        current_chunk: List[str] = [sentences[0]]
+
+        for i in range(1, len(sentences)):
+            start_prev = max(0, i - window_size)
+            end_prev = i
+            start_next = i
+            end_next = min(len(sentences), i + window_size)
+            
+            prev_window_emb = emb_tensor[start_prev:end_prev].mean(dim=0)
+            next_window_emb = emb_tensor[start_next:end_next].mean(dim=0)
+
+            similarity = F.cosine_similarity(
+                prev_window_emb.unsqueeze(0),
+                next_window_emb.unsqueeze(0),
+                dim=1
+            ).item()
+
+            dynamic_threshold = max(threshold * (1 - 0.03 * (end_next - start_prev)), 0.2)
+
+            if similarity < dynamic_threshold:
+                chunks.append(" ".join(current_chunk))
+                overlap_start = max(0, len(current_chunk) - self.chunk_overlap)
+                current_chunk = current_chunk[overlap_start:]
+                current_chunk.append(sentences[i])
+            else:
+                current_chunk.append(sentences[i])
+
+        if current_chunk:
+            chunks.append(" ".join(current_chunk))
+
+        return self.postprocess(chunks)