theroyallab/tabbyAPI
1
0
Fork 0
mirror of https://github.com/theroyallab/tabbyAPI.git synced 2026-04-29 10:42:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

migrate all yaml loaders to ruamel.yaml

Browse Source
This commit is contained in:
TerminalMan
2024-09-18 11:33:15 +01:00
parent 63634beb5e
commit 6c7542de9f
3 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
backends/exllamav2/model.py
View File

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

6
common/auth.py
View File

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

4
common/sampling.py
View File

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