squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix get_model

Browse Source
This commit is contained in:
Woosuk Kwon 2023-02-13 22:51:03 +00:00
parent 531e1c74e8
commit ee9442518d
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
cacheflow/worker/models/model_utils.py
View File

@ -8,6 +8,7 @@ MODEL_CLASSES = {
def get_model(model_name: str) -> nn.Module:
if model_name not in MODEL_CLASSES:
raise ValueError(f'Invalid model name: {model_name}')
return MODEL_CLASSES[model_name].from_pretrained(model_name)
for model_class, model in MODEL_CLASSES.items():
if model_class in model_name:
return model.from_pretrained(model_name)
raise ValueError(f'Invalid model name: {model_name}')