huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: LlamaTokenizerFast to AutoTokenizer at flash_mistral.py

Browse Source
This commit is contained in:
SeongBeomLEE 2024-03-11 13:27:09 +09:00
parent 7dbaf9e901
commit 2111ae1bd2
1 changed files with 17 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
server/text_generation_server/models/flash_mistral.py
View File

@ -6,7 +6,7 @@ import numpy as np
from dataclasses import dataclass from dataclasses import dataclass
from opentelemetry import trace from opentelemetry import trace
from transformers import PreTrainedTokenizerBase from transformers import PreTrainedTokenizerBase, AutoTokenizer
from transformers.models.llama import LlamaTokenizerFast from transformers.models.llama import LlamaTokenizerFast
from typing import Optional, Tuple, Type from typing import Optional, Tuple, Type
@ -317,13 +317,22 @@ class BaseFlashMistral(FlashCausalLM):
else: else:
raise NotImplementedError("FlashMistral is only available on GPU") raise NotImplementedError("FlashMistral is only available on GPU")
tokenizer = LlamaTokenizerFast.from_pretrained( try:
model_id, tokenizer = LlamaTokenizerFast.from_pretrained(
revision=revision, model_id,
padding_side="left", revision=revision,
truncation_side="left", padding_side="left",
trust_remote_code=trust_remote_code, truncation_side="left",
) trust_remote_code=trust_remote_code,
)
except Exception:
tokenizer = AutoTokenizer.from_pretrained(
model_id,
revision=revision,
padding_side="left",
truncation_side="left",
trust_remote_code=trust_remote_code,
)
config = config_cls.from_pretrained( config = config_cls.from_pretrained(
model_id, revision=revision, trust_remote_code=trust_remote_code model_id, revision=revision, trust_remote_code=trust_remote_code