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
This commit is contained in:
@@ -1,3 +1,4 @@
|
|||||||
|
from functools import cached_property
|
||||||
from typing import Any, Dict, List, Optional
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
from jinja2 import Template
|
from jinja2 import Template
|
||||||
@@ -29,7 +30,19 @@ class ChatTemplate:
|
|||||||
self.description = description
|
self.description = description
|
||||||
self.default_variables = default_variables or {}
|
self.default_variables = default_variables or {}
|
||||||
self.special_tokens = special_tokens 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
|
@classmethod
|
||||||
def from_string(
|
def from_string(
|
||||||
|
|||||||
@@ -164,7 +164,14 @@ class AutoTokenizer:
|
|||||||
- tokenizer.bos_token → returns string
|
- tokenizer.bos_token → returns string
|
||||||
- tokenizer.bos_token_id → returns corresponding integer ID
|
- tokenizer.bos_token_id → returns corresponding integer ID
|
||||||
- tokenizer.stop_ids → returns list of corresponding integer IDs for all special tokens
|
- 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
|
# Handle stop_ids - return IDs for all special tokens
|
||||||
if key == "stop_ids":
|
if key == "stop_ids":
|
||||||
stop_ids = []
|
stop_ids = []
|
||||||
|
|||||||
Reference in New Issue
Block a user