diff --git a/src/h3_blackwell_runtime/conditioning.py b/src/h3_blackwell_runtime/conditioning.py index b4d5f79..8103875 100644 --- a/src/h3_blackwell_runtime/conditioning.py +++ b/src/h3_blackwell_runtime/conditioning.py @@ -9,9 +9,13 @@ class H3PromptTokenizer: """Tokenize raw H3 prompt text without Qwen chat-template tokens.""" def __init__(self, tokenizer_dir: str | Path): - from transformers import AutoTokenizer + from transformers import Qwen2Tokenizer - self.tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_dir), local_files_only=True) + tokenizer_dir = Path(tokenizer_dir) + self.tokenizer = Qwen2Tokenizer( + vocab_file=str(tokenizer_dir / "vocab.json"), + merges_file=str(tokenizer_dir / "merges.txt"), + ) def __call__(self, prompt: str, *, device: torch.device | str = "cuda") -> torch.Tensor: if not prompt: