From 17127f8b3c716857966f0387964a1030ae7069f2 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sun, 19 Jul 2026 11:59:32 +0800
Subject: [PATCH] fix: make tokenizer picklable for spawn multiprocessing

- ChatTemplate: defer Jinja2 compilation to cached_property, exclude compiled template from __getstate__ (its dynamic root function has __module__=None and falls back to __main__, breaking pickle)
- AutoTokenizer: bypass __getattr__ for underscore-prefixed attrs to prevent infinite recursion during unpickle when __dict__ is empty
---
 astrai/tokenize/chat_template.py | 15 ++++++++++++++-
 astrai/tokenize/tokenizer.py     |  7 +++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/astrai/tokenize/chat_template.py b/astrai/tokenize/chat_template.py
index 77b2888..ea04ca1 100644
--- a/astrai/tokenize/chat_template.py
+++ b/astrai/tokenize/chat_template.py
@@ -1,3 +1,4 @@
+from functools import cached_property
 from typing import Any, Dict, List, Optional
 
 from jinja2 import Template
@@ -29,7 +30,19 @@ class ChatTemplate:
         self.description = description
         self.default_variables = default_variables or {}
         self.special_tokens = special_tokens or {}
-        self._compiled: Template = Template(template_str)
+
+    @cached_property
+    def _compiled(self) -> Template:
+        """Lazy-compiled Jinja2 template, cached on first access.
+
+        The compiled :class:`~jinja2.Template` holds a dynamically-generated
+        ``root`` render function whose ``__module__`` is ``None``; under
+        ``pickle`` it falls back to ``__main__`` and breaks ``spawn``-based
+        multiprocessing.  By deferring compilation to first access, the
+        default pickle protocol serialises only ``template_str``; each
+        worker rebuilds the cache on first render.
+        """
+        return Template(self.template_str)
 
     @classmethod
     def from_string(
diff --git a/astrai/tokenize/tokenizer.py b/astrai/tokenize/tokenizer.py
index bb883f0..83042db 100644
--- a/astrai/tokenize/tokenizer.py
+++ b/astrai/tokenize/tokenizer.py
@@ -164,7 +164,14 @@ class AutoTokenizer:
           - tokenizer.bos_token   → returns string
           - tokenizer.bos_token_id → returns corresponding integer ID
           - tokenizer.stop_ids → returns list of corresponding integer IDs for all special tokens
+
+        Internal/private attrs are not intercepted: during unpickle
+        ``__dict__`` is empty, so probing ``self._special_token_map``
+        would recurse infinitely.
         """
+        if key.startswith("_"):
+            raise AttributeError(key)
+
         # Handle stop_ids - return IDs for all special tokens
         if key == "stop_ids":
             stop_ids = []