"vscode:/vscode.git/clone" did not exist on "629b9ff288fc2e6b1a5bef7fb6e9e0467b28c906"
model_utils.py 340 Bytes
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
2
3
4
5
6
7
8
9
10
11
12
13
import torch.nn as nn

from cacheflow.worker.models.opt import OPTForCausalLM

MODEL_CLASSES = {
    'opt': OPTForCausalLM,
}


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)