refactor: 修改模型架构

pipeline/tokenize/chat_template.py

@@ -0,0 +1,120 @@
+"""
+Chat template module with Jinja2 rendering support.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from jinja2 import Template
+
+
+# Message type for chat messages
+type MessageType = Dict[str, Any]
+
+
+@dataclass
+class ChatTemplate:
+    """A chat template with Jinja2 rendering support.
+
+    Attributes:
+        name: Unique identifier for the template.
+        template_str: Jinja2 template string.
+        description: Optional description.
+        default_variables: Optional dictionary of default variable values
+            that will be passed to the template if not overridden during rendering.
+        special_tokens: Optional dictionary mapping token names to their string values.
+            These tokens are automatically added to the template variables.
+    """
+
+    name: str
+    template_str: str
+    description: str = ""
+    default_variables: Dict[str, Any] = field(default_factory=dict)
+    special_tokens: Dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def from_string(
+        cls,
+        template_str: str,
+        description: str = "",
+        default_variables: Optional[Dict[str, Any]] = None,
+        special_tokens: Optional[Dict[str, str]] = None,
+    ) -> "ChatTemplate":
+        """Create a ChatTemplate instance directly from a template string."""
+        return cls(
+            name="",  # empty name for ad-hoc templates
+            template_str=template_str,
+            description=description,
+            default_variables=default_variables or {},
+            special_tokens=special_tokens or {},
+        )
+
+    def render(
+        self,
+        messages: List[MessageType],
+        system_prompt: Optional[str] = None,
+        add_generation_prompt: bool = True,
+        **extra_variables: Any,
+    ) -> str:
+        """Render the template with given messages and variables.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'.
+            system_prompt: Optional system prompt string.
+            add_generation_prompt: Whether to add generation prompt after messages.
+            **extra_variables: Additional variables to pass to the template.
+                These override default_variables and special_tokens.
+
+        Returns:
+            Rendered prompt string.
+        """
+        # Merge default variables, special tokens, and extra variables
+        variables = {
+            **self.default_variables,
+            **self.special_tokens,
+            **extra_variables,
+        }
+        variables["messages"] = messages
+        variables["add_generation_prompt"] = add_generation_prompt
+        if system_prompt is not None:
+            variables["system_prompt"] = system_prompt
+
+        jinja_template = Template(self.template_str)
+        return jinja_template.render(**variables)
+
+
+# Default ChatML template
+DEFAULT_CHATML_TEMPLATE = """{% for message in messages %}{{ bos_token }}{{ message['role'] }}
+{{ message['content'] }}{{ eos_token }}{% endfor %}{% if add_generation_prompt %}{{ bos_token }}assistant
+{% endif %}"""
+
+
+# Pre-built template registry
+TEMPLATE_REGISTRY: Dict[str, ChatTemplate] = {}
+
+
+def register_chat_template(name: str, template: ChatTemplate) -> None:
+    """Register a chat template in the global registry.
+    
+    Args:
+        name: Name to register the template under
+        template: ChatTemplate instance
+    """
+    TEMPLATE_REGISTRY[name] = template
+
+
+def get_chat_template(name: str) -> ChatTemplate:
+    """Get a registered chat template.
+    
+    Args:
+        name: Template name
+        
+    Returns:
+        ChatTemplate instance
+        
+    Raises:
+        KeyError: If template not found
+    """
+    if name not in TEMPLATE_REGISTRY:
+        raise KeyError(f"Chat template '{name}' not found. Available: {list(TEMPLATE_REGISTRY.keys())}")
+    return TEMPLATE_REGISTRY[name]