diff --git a/generator/model.py b/generator/model.py index a6281e2..20818bb 100644 --- a/generator/model.py +++ b/generator/model.py @@ -10,7 +10,11 @@ def get_default_pipeline(model_name: str) -> Pipeline: model = AutoModelForCausalLM.from_pretrained( model_name, device_map="auto", torch_dtype="auto" ) - tokenizer = AutoTokenizer.from_pretrained(model_name) + + try: + tokenizer = AutoTokenizer.from_pretrained(model_name) + except: + tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast = False) pipe = tf_pipe(task="text-generation", model=model, tokenizer=tokenizer) @@ -19,7 +23,11 @@ def get_default_pipeline(model_name: str) -> Pipeline: def get_onnx_pipeline(model_name: str) -> Pipeline: model = ORTModelForCausalLM.from_pretrained(model_name) - tokenizer = AutoTokenizer.from_pretrained(model_name) + + try: + tokenizer = AutoTokenizer.from_pretrained(model_name) + except: + tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast = False) pipe = opt_pipe( task="text-generation", model=model, tokenizer=tokenizer, accelerator="ort" @@ -32,7 +40,11 @@ def get_bettertransformer_pipeline(model_name: str) -> Pipeline: model = AutoModelForCausalLM.from_pretrained( model_name, device_map="auto", torch_dtype="auto" ) - tokenizer = AutoTokenizer.from_pretrained(model_name) + + try: + tokenizer = AutoTokenizer.from_pretrained(model_name) + except: + tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast = False) pipe = opt_pipe( task="text-generation",