Tree: Use safe loader for YAML

Loaders that read use a safe type while loaders that write use both
round-trip and safe options.

Also don't create module-level parsers where they're not needed.

Signed-off-by: kingbri <bdashore3@proton.me>
This commit is contained in:
kingbri
2024-09-18 12:12:49 -04:00
parent 6c7542de9f
commit 24ea85b3c5
4 changed files with 18 additions and 12 deletions

View File

@@ -56,8 +56,6 @@ 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."""
@@ -381,7 +379,10 @@ class ExllamaV2Container:
override_config_path, "r", encoding="utf8"
) as override_config_file:
contents = await override_config_file.read()
override_args = unwrap(yaml.safe_load(contents), {})
# Create a temporary YAML parser
yaml = YAML(typ="safe")
override_args = unwrap(yaml.load(contents), {})
# Merge draft overrides beforehand
draft_override_args = unwrap(override_args.get("draft"), {})