Add model_utils

cacheflow/worker/models/model_utils.py

@@ -0,0 +1,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)