2023-02-13 17:36:12 +08:00
|
|
|
import torch.nn as nn
|
|
|
|
|
|
2023-02-23 02:08:25 +08:00
|
|
|
from cacheflow.models.opt import OPTForCausalLM
|
2023-02-13 17:36:12 +08:00
|
|
|
|
|
|
|
|
MODEL_CLASSES = {
|
|
|
|
|
'opt': OPTForCausalLM,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_model(model_name: str) -> nn.Module:
|
2023-02-14 06:51:03 +08:00
|
|
|
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}')
|
