Update crs_arena/utils.py

crs_arena/utils.py

@@ -27,7 +27,37 @@ HF_API = HfApi(token=st.secrets["hf_token"])
 
 
 @st.cache_resource(
-    show_spinner="Loading CRS...", ttl=timedelta(days=3)
+    show_spinner="<h2 style='text-align: center; top:50%; left:50%; transform: translate(-50%, -50%);'>The fighters are warming up... :robot: :punch: :gun: :boom:</h2>",
+    ttl=timedelta(days=3),
+)
+def get_crs_model(model_name: str, model_config_file: str) -> CRSModel:
+    """Returns a CRS model.
+
+    Args:
+        model_name: Model name.
+        model_config_file: Model configuration file.
+
+    Raises:
+        FileNotFoundError: If model configuration file is not found.
+
+    Returns:
+        CRS model.
+    """
+    logging.debug(f"Loading CRS model {model_name}.")
+    if not os.path.exists(model_config_file):
+        raise FileNotFoundError(
+            f"Model configuration file {model_config_file} not found."
+        )
+
+    model_args = yaml.safe_load(open(model_config_file, "r"))
+
+    if "chatgpt" in model_name:
+        openai.api_key = st.secrets["openai_api_key"]
+
+    # Extract crs model from name
+    name = model_name.split("_")[0]
+
+    return CRSModel(name, **model_args), ttl=timedelta(days=3)
 )
 def get_crs_model(model_name: str, model_config_file: str) -> CRSModel:
     """Returns a CRS model.