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 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(
|
||||
|
||||
@@ -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 = []
|
||||
|
||||
Reference in New Issue
Block a user