Load Qwen tokenizer from local vocab
This commit is contained in:
parent
b188002199
commit
7e2062ee38
1 changed files with 6 additions and 2 deletions
|
|
@ -9,9 +9,13 @@ class H3PromptTokenizer:
|
||||||
"""Tokenize raw H3 prompt text without Qwen chat-template tokens."""
|
"""Tokenize raw H3 prompt text without Qwen chat-template tokens."""
|
||||||
|
|
||||||
def __init__(self, tokenizer_dir: str | Path):
|
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:
|
def __call__(self, prompt: str, *, device: torch.device | str = "cuda") -> torch.Tensor:
|
||||||
if not prompt:
|
if not prompt:
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue