migrate all yaml loaders to ruamel.yaml

backends/exllamav2/model.py

@@ -30,7 +30,7 @@ from itertools import zip_longest
 from loguru import logger
 from typing import List, Optional, Union
 
-import yaml
+from ruamel.yaml import YAML
 
 from backends.exllamav2.grammar import (
     ExLlamaV2Grammar,
@@ -56,6 +56,8 @@ from common.templating import (
 from common.transformers_utils import GenerationConfig, HuggingFaceConfig
 from common.utils import coalesce, unwrap
 
+yaml = YAML()
+
 
 class ExllamaV2Container:
     """The model container class for ExLlamaV2 models."""

common/auth.py

@@ -5,7 +5,7 @@ application, it should be fine.
 
 import aiofiles
 import secrets
-import yaml
+from ruamel.yaml import YAML
 from fastapi import Header, HTTPException, Request
 from pydantic import BaseModel
 from loguru import logger
@@ -13,6 +13,8 @@ from typing import Optional
 
 from common.utils import coalesce
 
+yaml = YAML()
+
 
 class AuthKeys(BaseModel):
     """
@@ -60,7 +62,7 @@ async def load_auth_keys(disable_from_config: bool):
     try:
         async with aiofiles.open("api_tokens.yml", "r", encoding="utf8") as auth_file:
             contents = await auth_file.read()
-            auth_keys_dict = yaml.safe_load(contents)
+            auth_keys_dict = yaml.load(contents)
             AUTH_KEYS = AuthKeys.model_validate(auth_keys_dict)
     except FileNotFoundError:
         new_auth_keys = AuthKeys(

common/sampling.py

@@ -3,7 +3,7 @@
 import aiofiles
 import json
 import pathlib
-import yaml
+from ruamel.yaml import YAML
 from copy import deepcopy
 from loguru import logger
 from pydantic import AliasChoices, BaseModel, Field
@@ -11,6 +11,8 @@ from typing import Dict, List, Optional, Union
 
 from common.utils import unwrap, prune_dict
 
+yaml = YAML()
+
 
 # Common class for sampler params
 class BaseSamplerRequest(BaseModel):