fix: LlamaTokenizerFast to AutoTokenizer at flash_llama.py

This commit is contained in:
Dong Shin 2023-07-16 21:09:02 +09:00 committed by GitHub
parent a2cf1bdb2f
commit abe4e4b1cc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -2,8 +2,8 @@ import torch
import torch.distributed import torch.distributed
from opentelemetry import trace from opentelemetry import trace
from transformers import AutoConfig from transformers import AutoConfig, AutoTokenizer
from transformers.models.llama import LlamaTokenizer, LlamaTokenizerFast from transformers.models.llama import LlamaTokenizer
from typing import Optional from typing import Optional
from text_generation_server.models import FlashCausalLM from text_generation_server.models import FlashCausalLM
@ -44,7 +44,7 @@ class FlashLlama(FlashCausalLM):
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
) )
except Exception: except Exception:
tokenizer = LlamaTokenizerFast.from_pretrained( tokenizer = AutoTokenizer.from_pretrained(
model_id, model_id,
revision=revision, revision=revision,
padding_side="left", padding_side="left",