Add model_utils

cacheflow/worker/models/__init__.py

@@ -1,5 +1,7 @@
-from cacheflow.worker.models.opt import OPTForCausalLM
+from cacheflow.worker.models.model_utils import get_model
+
 
 __all__ = [
-    'OPTForCausalLM',
+    'get_model',
+    
 ]

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)