nodes/alpertunga-bile-prompt-generator-comfyui
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

added try/except for tokenizer inits

Browse Source
This commit is contained in:
alpertunga-bile
2024-03-05 13:46:24 +00:00
committed by GitHub
parent 536a4e8b31
commit b7b7ee8083
1 changed files with 15 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
generator/model.py
+15 -3
View File
@@ -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",