Skip to content

Commit

Permalink
fix falcon initialization during benchmarking (#362)
Browse files Browse the repository at this point in the history
  • Loading branch information
eaidova authored Apr 15, 2024
1 parent e4238b7 commit 534963a
Showing 1 changed file with 10 additions and 1 deletion.
11 changes: 10 additions & 1 deletion llm_bench/python/utils/ov_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,7 @@ def create_text_gen_model(model_path, device, **kwargs):
if not model_path_existed:
raise RuntimeError(f'==Failure ==: model path:{model_path} does not exist')
else:
if model_type in ['mpt', 'falcon', 'replit', 'codegen2', 'chatglm']:
if model_type in ['replit', 'codegen2', 'chatglm', 'mpt']:
start = time.perf_counter()
ov_model = model_class.from_pretrained(
model_path,
Expand All @@ -151,6 +151,15 @@ def create_text_gen_model(model_path, device, **kwargs):
stateful=kwargs.get("stateful", None)
)
end = time.perf_counter()
elif model_type in ['falcon']:
ov_model = model_class.from_pretrained(
model_path,
device=device,
ov_config=ov_config,
stateful=kwargs.get("stateful", None),
trust_remote_code=False
)
end = time.perf_counter()
else:
start = time.perf_counter()
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
Expand Down

0 comments on commit 534963a

Please sign in to comment.