orbit-torch 0.0.4a1__py3-none-any.whl → 0.1.0b1__py3-none-any.whl

orbit/kit/interface.py

@@ -0,0 +1,154 @@
+from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.live import Live
+from rich.prompt import Prompt
+
+from .wrapper import AutoRegressiveWrapper
+
+class ChatInterface:
+    '''
+    一个用于命令行实时交互的聊天接口类
+
+    Attributes:
+        model: 语言模型实例
+        tokenizer: 分词器实例
+        device: 模型所在的设备
+    '''
+
+    def __init__(self, model=None, tokenizer=None, model_id=None, device='auto', dtype='auto', model_role='assistant'):
+        '''
+        初始化聊天接口
+
+        Args:
+            model: 预加载的模型实例。如果为 None，则需要提供 model_id
+            tokenizer: 预加载的分词器实例。如果为 None，则需要提供 model_id
+            model_id: 模型的 HuggingFace ID 或本地路径
+            device: 设备设置，默认为 'auto'
+            dtype: 模型的权重精度，默认为 'auto'
+            model_role: 模型回复的角色名称，默认为 'assistant'
+        '''
+        self.console = Console()
+        self.model_role = model_role
+        if model is not None and tokenizer is not None:
+            self.model = model
+            self.tokenizer = tokenizer
+        elif model_id is not None:
+            self._load_model(model_id, device, dtype)
+        else:
+            raise ValueError('必须提供 (model 和 tokenizer) 或 model_id')
+
+        self.device = self.model.device
+
+        if not hasattr(self.model, 'generate'):
+            self.model = AutoRegressiveWrapper(self.model)
+
+    def _load_model(self, model_id, device, dtype):
+        '''
+        从指定的 model_id 加载模型和分词器
+
+        Args:
+            model_id: 模型的 HuggingFace ID 或本地路径
+            device: 设备设置
+            dtype: 模型精度
+        '''
+        with self.console.status(f'[bold green]正在加载模型: {model_id} ...[/bold green]'):
+            self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                dtype=dtype,
+                device_map=device
+            )
+        self.console.print(f'[bold green]模型 {model_id} 加载完成！[/bold green]')
+
+    def stream_chat(
+        self,
+        messages: list,
+        max_new_tokens: int = 512,
+        temperature: float = 0.7,
+        top_k: int = 50,
+        top_p: float = 0.9,
+        repetition_penalty: float = 1.0,
+        do_sample: bool = True
+    ):
+        '''
+        流式生成对话响应
+
+        Args:
+            messages (list): 符合 ChatML 格式的消息列表
+            max_new_tokens (int): 最大新生成的 token 数量
+            temperature (float): 生成温度
+            top_k (int): Top-k 采样值
+            top_p (float): Top-p 采样值
+            repetition_penalty (float): 重复惩罚系数
+            do_sample (bool): 是否使用采样
+
+        Yields:
+            str: 生成的新文本片段
+        '''
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        model_inputs = self.tokenizer([text], return_tensors='pt').to(self.device)
+
+        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+        generation_kwargs = dict(
+            model_inputs,
+            streamer=streamer,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            do_sample=do_sample,
+            eos_token_id=self.tokenizer.eos_token_id
+        )
+
+        thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        for new_text in streamer:
+            yield new_text
+
+    def interact(self):
+        '''
+        启动命令行实时交互会话
+        '''
+        self.console.print(Panel(
+            '[bold]聊天接口已就绪[/bold]\n输入 [red]"exit"[/red] 或 [red]"quit"[/red] 退出',
+            title='[bold blue]Orbit Chat[/bold blue]',
+            border_style='blue',
+            expand=False
+        ))
+        history = []
+        while True:
+            try:
+                self.console.print()
+                user_input = Prompt.ask('[bold green]User[/bold green]', console=self.console)
+                if user_input.lower() in ['exit', 'quit']:
+                    break
+                if not user_input.strip():
+                    continue
+
+                history.append({'role': 'user', 'content': user_input})
+                
+                self.console.print('[bold purple]Model:[/bold purple]\n')
+                full_response = ''
+                
+                with Live(Markdown(""), console=self.console, refresh_per_second=12) as live:
+                    for chunk in self.stream_chat(history):
+                        full_response += chunk
+                        live.update(Markdown(full_response))
+                
+                history.append({'role': self.model_role, 'content': full_response})
+
+            except KeyboardInterrupt:
+                self.console.print('\n[bold red][会话已中断][/bold red]')
+                break
+        
+        self.console.print('[bold blue][生成结束][/bold blue]')

orbit/kit/wrapper.py

@@ -0,0 +1,157 @@
+import torch
+import torch.nn as nn
+import inspect
+from typing import Optional, Any, List, Union
+
+
+class AutoRegressiveWrapper:
+    '''
+    一个将普通 torch.nn.Module 包装为兼容 transformers generate 接口的包装类
+
+    Attributes:
+        model (nn.Module): 原始模型实例
+        device (torch.device): 模型所在的设备
+        accepts_attention_mask (bool): 模型是否接受 attention_mask 参数
+        accepts_mask (bool): 模型是否接受 mask 参数
+    '''
+
+    def __init__(self, model: nn.Module):
+        '''
+        初始化包装器
+
+        Args:
+            model (nn.Module): 只有 forward 方法的自定义模型
+        '''
+        self.model = model
+        try:
+            self.device = next(model.parameters()).device
+        except StopIteration:
+            self.device = torch.device('cpu')
+
+        sig = inspect.signature(model.forward)
+        params = sig.parameters
+        self.accepts_attention_mask = 'attention_mask' in params
+        self.accepts_mask = 'mask' in params
+        
+        self.model.eval()
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 512,
+        temperature: float = 1.0,
+        top_k: int = 50,
+        top_p: float = 0.9,
+        repetition_penalty: float = 1.0,
+        do_sample: bool = True,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        streamer: Optional[Any] = None,
+        **kwargs
+    ) -> torch.Tensor:
+        '''
+        自回归生成循环，支持多种采样策略并兼容 TextIteratorStreamer
+
+        Args:
+            input_ids (torch.Tensor): 输入的 token ID 序列 [batch, seq_len]
+            max_new_tokens (int): 最大新生成的 token 数量
+            temperature (float): 采样温度
+            top_k (int): Top-k 采样的 k 值
+            top_p (float): Top-p (Nucleus) 采样的 p 值
+            repetition_penalty (float): 重复惩罚系数
+            do_sample (bool): 是否使用采样
+            eos_token_id (Optional[Union[int, List[int]]]): 终止 token ID
+            streamer (Optional[Any]): transformers 库的 streamer 实例
+            **kwargs: 忽略其他 transformers 相关的参数
+
+        Returns:
+            torch.Tensor: 包含生成内容的完整序列
+        '''
+        curr_input_ids = input_ids.to(self.device)
+        batch_size = curr_input_ids.shape[0]
+        
+        if isinstance(eos_token_id, int):
+            eos_token_id = [eos_token_id]
+        
+        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=self.device)
+
+        for _ in range(max_new_tokens):
+            model_inputs = {'input_ids': curr_input_ids}
+
+            if self.accepts_attention_mask:
+                model_inputs['attention_mask'] = torch.ones_like(curr_input_ids)
+            elif self.accepts_mask:
+                model_inputs['mask'] = torch.ones_like(curr_input_ids)
+
+            outputs = self.model(**model_inputs)
+
+            if isinstance(outputs, (tuple, list)):
+                logits = outputs[0]
+            else:
+                logits = outputs
+
+            next_token_logits = logits[:, -1, :]
+
+            if repetition_penalty != 1.0:
+                for i in range(batch_size):
+                    for previous_token in set(curr_input_ids[i].tolist()):
+                        if next_token_logits[i, previous_token] < 0:
+                            next_token_logits[i, previous_token] *= repetition_penalty
+                        else:
+                            next_token_logits[i, previous_token] /= repetition_penalty
+
+            if do_sample:
+                if temperature != 1.0:
+                    next_token_logits = next_token_logits / temperature
+
+                if top_k > 0:
+                    indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                    next_token_logits[indices_to_remove] = float('-inf')
+
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                    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 = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                    next_token_logits[indices_to_remove] = float('-inf')
+
+                probs = torch.softmax(next_token_logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+            
+            curr_input_ids = torch.cat([curr_input_ids, next_token], dim=-1)
+
+            if streamer is not None:
+                if unfinished_sequences[0] == 1:
+                    streamer.put(next_token.cpu())
+
+            if eos_token_id is not None:
+                for token_id in eos_token_id:
+                    unfinished_sequences = unfinished_sequences.mul(
+                        next_token.tile(1, 1).ne(token_id).all(dim=-1).long()
+                    )
+
+            if unfinished_sequences.max() == 0:
+                break
+
+        if streamer is not None:
+            streamer.end()
+
+        return curr_input_ids
+
+    def __getattr__(self, name: str) -> Any:
+        '''
+        将未定义的属性访问转发给原始模型
+
+        Args:
+            name (str): 属性名称
+
+        Returns:
+            Any: 原始模型的属性
+        '''
+        return getattr(self.model, name)

orbit/model/__init__.py

@@ -0,0 +1,5 @@
+from .base import BaseBlock
+from .registry import (
+    register_model, build_model, list_models, get_model_class
+)
+from .config import ModelConfig

orbit/model/base.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+from typing import Union, List, Optional, Iterable
+
+from orbit.utils import (
+    auto_initialize,
+    freeze_layers,
+    unfreeze_layers,
+    count_params,
+    save_model,
+    load_model,
+)
+
+
+class BaseBlock(nn.Module):
+    ''' 基础模型块，提供通用的模型功能。
+
+    继承自 nn.Module，包含参数统计、梯度检查点、冻结/解冻层、保存/加载模型等功能。
+    '''
+
+    def __init__(self):
+        ''' 初始化 BaseBlock。 '''
+        super(BaseBlock, self).__init__()
+
+        self.gradient_checkpointing: bool = False
+    
+    @property
+    def device(self):
+        ''' 获取模型所在的设备。
+
+        Returns:
+            torch.device: 模型参数所在的设备。如果没有参数，则返回 cpu。
+        '''
+        try:
+            return next(self.parameters()).device
+        except StopIteration:
+            return torch.device('cpu')
+
+    def _init_weights(self, model: Union[nn.Module, 'BaseBlock', nn.Parameter, torch.Tensor]):
+        ''' 初始化模型权重。
+
+        Args:
+            model (Union[nn.Module, 'BaseBlock', nn.Parameter, torch.Tensor]): 需要初始化的模型、层或张量。
+        '''
+        auto_initialize(model=model, verbose=False)
+    
+    def set_checkpoint(self, value: bool):
+        ''' 设置是否启用梯度检查点。
+
+        Args:
+            value (bool): 是否启用梯度检查点。
+        '''
+        self.gradient_checkpointing = value
+        for model in self.modules():
+            if isinstance(model, BaseBlock) and model is not self:
+                model.gradient_checkpointing = value
+
+    def count_params(self, trainable_only=False):
+        ''' 统计模型参数数量。
+
+        Args:
+            trainable_only (bool, optional): 是否只统计可训练参数。默认为 False。
+
+        Returns:
+            int: 参数数量。
+        '''
+        if trainable_only:
+            return count_params(self).count
+        
+        return count_params(self, mode='all').count
+
+    def checkpoint(self, function, *args, **kwargs):
+        ''' 应用梯度检查点。
+
+        如果启用了梯度检查点且处于训练模式，则使用 torch.utils.checkpoint.checkpoint。
+        否则直接调用函数。
+
+        Args:
+            function (Callable): 要执行的函数。
+            *args: 传递给函数的位置参数。
+            **kwargs: 传递给函数的关键字参数。
+
+        Returns:
+            Any: 函数的返回值。
+        '''
+        if self.gradient_checkpointing and self.training:
+            return torch.utils.checkpoint.checkpoint(function, *args, use_reentrant=False, **kwargs)
+        else:
+            return function(*args, **kwargs)
+
+    def freeze(self, targets: Optional[Union[str, List[str]]] = None):
+        ''' 冻结指定层的参数。
+
+        Args:
+            targets (Optional[Union[str, List[str]]], optional): 要冻结的层名称或名称列表。
+                如果为 None，则冻结所有层。默认为 None。
+        '''
+        freeze_layers(self, targets)
+
+    def unfreeze(self, targets: Optional[Union[str, List[str]]] = None):
+        ''' 解冻指定层的参数。
+
+        Args:
+            targets (Optional[Union[str, List[str]]], optional): 要解冻的层名称或名称列表。
+                如果为 None，则解冻所有层。默认为 None。
+        '''
+        unfreeze_layers(self, targets)
+
+    def save_pretrained(self, file_path: str):
+        ''' 保存模型权重到文件。
+
+        Args:
+            file_path (str): 保存路径。
+        '''
+        save_model(self, file_path)
+
+    def load_pretrained(self, file_path: str, strict: bool = True, map_location: Union[str, torch.device] = 'cpu'):
+        ''' 从文件加载模型权重。
+
+        Args:
+            file_path (str): 权重文件路径。
+            strict (bool, optional): 是否严格匹配键值。默认为 True。
+            map_location (Union[str, torch.device], optional): 映射位置。默认为 'cpu'。
+        '''
+        load_model(self, file_path, strict, map_location)

orbit/model/block/__init__.py

@@ -0,0 +1,34 @@
+from .embedding  import (
+    RotaryPositionalEmbedding,
+    SinusoidalPositionalEmbedding,
+    MRoPEInterleavedEmbedding
+)
+from .attention import (
+    MultiHeadAttention, apply_attention, AttentionOutput,
+    SpatialMultiHeadAttention
+)
+from .codebook import (
+    LFQ, QuantizerOutput
+)
+from .fusion import (
+    LowRankFusion, GatedMultimodalUnit, DiffusionMapsFusion, CompactMultimodalPooling
+)
+from .mlp  import MLP
+from .moe  import MoE
+from .tcn  import TCN
+from .bio  import (
+    HebianLayer, PredictiveCodingLayer, PredictiveCodingOutput,
+    PredictiveCodingBlock
+)
+from .film import FiLM, FiLMOutput
+from .gate import (
+    SigmoidGate, TanhGate, SoftmaxGate, GLUGate,
+    TopKGate, TopKGateOutput
+)
+from .conv import (
+    CausalConv1d, calculate_causal_layer, ConvBlock,
+    DepthwiseSeparableConv, ResBasicBlock
+)
+from .lora import (
+    LinearLoRA, Conv2dLoRA, Conv1dLoRA, EmbeddingLoRA
+)