config is now backed by pydantic (WIP)

- add models for config options
- add function to regenerate config.yml
- replace references to config with pydantic compatible references
- remove unnecessary unwrap() statements

TODO:

- auto generate env vars
- auto generate argparse
- test loading a model

common/config_models.py

@@ -0,0 +1,248 @@
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from typing import List, Optional, Union, get_type_hints
+
+from common.utils import unwrap
+
+
+class config_config_model(BaseModel):
+    config: Optional[str] = Field(
+        None, description="Path to an overriding config.yml file"
+    )
+
+
+class network_config_model(BaseModel):
+    host: Optional[str] = Field("127.0.0.1", description="The IP to host on")
+    port: Optional[int] = Field(5000, description="The port to host on")
+    disable_auth: Optional[bool] = Field(
+        False, description="Disable HTTP token authentication with requests"
+    )
+    send_tracebacks: Optional[bool] = Field(
+        False, description="Decide whether to send error tracebacks over the API"
+    )
+    api_servers: Optional[List[str]] = Field(
+        [
+            "OAI",
+        ],
+        description="API servers to enable. Options: (OAI, Kobold)",
+    )
+
+
+class logging_config_model(BaseModel):
+    log_prompt: Optional[bool] = Field(False, description="Enable prompt logging")
+    log_generation_params: Optional[bool] = Field(
+        False, description="Enable generation parameter logging"
+    )
+    log_requests: Optional[bool] = Field(False, description="Enable request logging")
+
+
+class model_config_model(BaseModel):
+    model_dir: str = Field(
+        "models",
+        description="Overrides the directory to look for models (default: models). Windows users, do NOT put this path in quotes.",
+    )
+    use_dummy_models: Optional[bool] = Field(
+        False,
+        description="Sends dummy model names when the models endpoint is queried. Enable this if looking for specific OAI models.",
+    )
+    model_name: Optional[str] = Field(
+        None,
+        description="An initial model to load. Make sure the model is located in the model directory! REQUIRED: This must be filled out to load a model on startup.",
+    )
+    use_as_default: List[str] = Field(
+        default_factory=list,
+        description="Names of args to use as a default fallback for API load requests (default: []). Example: ['max_seq_len', 'cache_mode']",
+    )
+    max_seq_len: Optional[int] = Field(
+        None,
+        description="Max sequence length. Fetched from the model's base sequence length in config.json by default.",
+    )
+    override_base_seq_len: Optional[int] = Field(
+        None,
+        description="Overrides base model context length. WARNING: Only use this if the model's base sequence length is incorrect.",
+    )
+    tensor_parallel: Optional[bool] = Field(
+        False,
+        description="Load model with tensor parallelism. Fallback to autosplit if GPU split isn't provided.",
+    )
+    gpu_split_auto: Optional[bool] = Field(
+        True,
+        description="Automatically allocate resources to GPUs (default: True). Not parsed for single GPU users.",
+    )
+    autosplit_reserve: List[int] = Field(
+        [96],
+        description="Reserve VRAM used for autosplit loading (default: 96 MB on GPU 0). Represented as an array of MB per GPU.",
+    )
+    gpu_split: List[float] = Field(
+        default_factory=list,
+        description="An integer array of GBs of VRAM to split between GPUs (default: []). Used with tensor parallelism.",
+    )
+    rope_scale: Optional[float] = Field(
+        1.0,
+        description="Rope scale (default: 1.0). Same as compress_pos_emb. Only use if the model was trained on long context with rope.",
+    )
+    rope_alpha: Optional[Union[float, str]] = Field(
+        1.0,
+        description="Rope alpha (default: 1.0). Same as alpha_value. Set to 'auto' to auto-calculate.",
+    )
+    cache_mode: Optional[str] = Field(
+        "FP16",
+        description="Enable different cache modes for VRAM savings (default: FP16). Possible values: FP16, Q8, Q6, Q4.",
+    )
+    cache_size: Optional[int] = Field(
+        None,
+        description="Size of the prompt cache to allocate (default: max_seq_len). Must be a multiple of 256.",
+    )
+    chunk_size: Optional[int] = Field(
+        2048,
+        description="Chunk size for prompt ingestion (default: 2048). A lower value reduces VRAM usage but decreases ingestion speed.",
+    )
+    max_batch_size: Optional[int] = Field(
+        None,
+        description="Set the maximum number of prompts to process at one time (default: None/Automatic). Automatically calculated if left blank.",
+    )
+    prompt_template: Optional[str] = Field(
+        None,
+        description="Set the prompt template for this model. If empty, attempts to look for the model's chat template.",
+    )
+    num_experts_per_token: Optional[int] = Field(
+        None,
+        description="Number of experts to use per token. Fetched from the model's config.json. For MoE models only.",
+    )
+    fasttensors: Optional[bool] = Field(
+        False,
+        description="Enables fasttensors to possibly increase model loading speeds (default: False).",
+    )
+
+
+class draft_model_config_model(BaseModel):
+    draft_model_dir: Optional[str] = Field(
+        "models",
+        description="Overrides the directory to look for draft models (default: models)",
+    )
+    draft_model_name: Optional[str] = Field(
+        None,
+        description="An initial draft model to load. Ensure the model is in the model directory.",
+    )
+    draft_rope_scale: Optional[float] = Field(
+        1.0,
+        description="Rope scale for draft models (default: 1.0). Same as compress_pos_emb. Use if the draft model was trained on long context with rope.",
+    )
+    draft_rope_alpha: Optional[float] = Field(
+        None,
+        description="Rope alpha for draft models (default: None). Same as alpha_value. Leave blank to auto-calculate the alpha value.",
+    )
+    draft_cache_mode: Optional[str] = Field(
+        "FP16",
+        description="Cache mode for draft models to save VRAM (default: FP16). Possible values: FP16, Q8, Q6, Q4.",
+    )
+
+
+class lora_instance_model(BaseModel):
+    name: str = Field(..., description="Name of the LoRA model")
+    scaling: float = Field(
+        1.0, description="Scaling factor for the LoRA model (default: 1.0)"
+    )
+
+
+class lora_config_model(BaseModel):
+    lora_dir: Optional[str] = Field(
+        "loras", description="Directory to look for LoRAs (default: 'loras')"
+    )
+    loras: Optional[List[lora_instance_model]] = Field(
+        None,
+        description="List of LoRAs to load and associated scaling factors (default scaling: 1.0)",
+    )
+
+
+class sampling_config_model(BaseModel):
+    override_preset: Optional[str] = Field(
+        None, description="Select a sampler override preset"
+    )
+
+
+class developer_config_model(BaseModel):
+    unsafe_launch: Optional[bool] = Field(
+        False, description="Skip Exllamav2 version check"
+    )
+    disable_request_streaming: Optional[bool] = Field(
+        False, description="Disables API request streaming"
+    )
+    cuda_malloc_backend: Optional[bool] = Field(
+        False, description="Runs with the pytorch CUDA malloc backend"
+    )
+    uvloop: Optional[bool] = Field(
+        False, description="Run asyncio using Uvloop or Winloop"
+    )
+    realtime_process_priority: Optional[bool] = Field(
+        False,
+        description="Set process to use a higher priority For realtime process priority, run as administrator or sudo Otherwise, the priority will be set to high",
+    )
+
+
+class embeddings_config_model(BaseModel):
+    embedding_model_dir: Optional[str] = Field(
+        "models",
+        description="Overrides directory to look for embedding models (default: models)",
+    )
+    embeddings_device: Optional[str] = Field(
+        "cpu",
+        description="Device to load embedding models on (default: cpu). Possible values: cpu, auto, cuda. If using an AMD GPU, set this value to 'cuda'.",
+    )
+    embedding_model_name: Optional[str] = Field(
+        None, description="The embeddings model to load"
+    )
+
+
+class tabby_config_model(BaseModel):
+    config: config_config_model = Field(default_factory=config_config_model)
+    network: network_config_model = Field(default_factory=network_config_model)
+    logging: logging_config_model = Field(default_factory=logging_config_model)
+    model: model_config_model = Field(default_factory=model_config_model)
+    draft_model: draft_model_config_model = Field(
+        default_factory=draft_model_config_model
+    )
+    lora: lora_config_model = Field(default_factory=lora_config_model)
+    sampling: sampling_config_model = Field(default_factory=sampling_config_model)
+    developer: developer_config_model = Field(default_factory=developer_config_model)
+    embeddings: embeddings_config_model = Field(default_factory=embeddings_config_model)
+
+    @model_validator(mode="before")
+    def set_defaults(cls, values):
+        for field_name, field_value in values.items():
+            if field_value is None:
+                default_instance = cls.__annotations__[field_name]().dict()
+                values[field_name] = cls.__annotations__[field_name](**default_instance)
+        return values
+
+    model_config = ConfigDict(validate_assignment=True)
+
+
+def generate_config_file(filename="config_sample.yml", indentation=2):
+    schema = tabby_config_model.model_json_schema()
+
+    def dump_def(id: str, indent=2):
+        yaml = ""
+        indent = " " * indentation * indent
+        id = id.split("/")[-1]
+
+        section = schema["$defs"][id]["properties"]
+        for property in section.keys():  # get type
+            comment = section[property]["description"]
+            yaml += f"{indent}# {comment}\n"
+
+            value = unwrap(section[property].get("default"), "")
+            yaml += f"{indent}{property}: {value}\n\n"
+
+        return yaml + "\n"
+
+    yaml = ""
+    for section in schema["properties"].keys():
+        yaml += f"{section}:\n"
+        yaml += dump_def(schema["properties"][section]["$ref"])
+        yaml += "\n"
+
+    with open(filename, "w") as f:
+        f.write(yaml)
+
+
+# generate_config_file("test.yml")

common/downloader.py

@@ -76,9 +76,9 @@ def _get_download_folder(repo_id: str, repo_type: str, folder_name: Optional[str
     """Gets the download folder for the repo."""
 
     if repo_type == "lora":
-        download_path = pathlib.Path(config.lora.get("lora_dir") or "loras")
+        download_path = pathlib.Path(config.lora.lora_dir)
     else:
-        download_path = pathlib.Path(config.model.get("model_dir") or "models")
+        download_path = pathlib.Path(config.model.model_dir)
 
     download_path = download_path / (folder_name or repo_id.split("/")[-1])
     return download_path

common/gen_logging.py

@@ -6,37 +6,19 @@ from pydantic import BaseModel
 from loguru import logger
 from typing import Dict, Optional
 
-
-class GenLogPreferences(BaseModel):
-    """Logging preference config."""
-
-    prompt: bool = False
-    generation_params: bool = False
-
+from common.tabby_config import config
 
 # Global logging preferences constant
-PREFERENCES = GenLogPreferences()
-
-
-def update_from_dict(options_dict: Dict[str, bool]):
-    """Wrapper to set the logging config for generations"""
-    global PREFERENCES
-
-    # Force bools on the dict
-    for value in options_dict.values():
-        if value is None:
-            value = False
-
-    PREFERENCES = GenLogPreferences.model_validate(options_dict)
+PREFERENCES = config.logging
 
 
 def broadcast_status():
     """Broadcasts the current logging status"""
     enabled = []
-    if PREFERENCES.prompt:
+    if PREFERENCES.log_prompt:
         enabled.append("prompts")
 
-    if PREFERENCES.generation_params:
+    if PREFERENCES.log_generation_params:
         enabled.append("generation params")
 
     if len(enabled) > 0:
@@ -47,13 +29,13 @@ def broadcast_status():
 
 def log_generation_params(**kwargs):
     """Logs generation parameters to console."""
-    if PREFERENCES.generation_params:
+    if PREFERENCES.log_generation_params:
         logger.info(f"Generation options: {kwargs}\n")
 
 
 def log_prompt(prompt: str, request_id: str, negative_prompt: Optional[str]):
     """Logs the prompt to console."""
-    if PREFERENCES.prompt:
+    if PREFERENCES.log_prompt:
         formatted_prompt = "\n" + prompt
         logger.info(
             f"Prompt (ID: {request_id}): {formatted_prompt if prompt else 'Empty'}\n"
@@ -66,7 +48,7 @@ def log_prompt(prompt: str, request_id: str, negative_prompt: Optional[str]):
 
 def log_response(request_id: str, response: str):
     """Logs the response to console."""
-    if PREFERENCES.prompt:
+    if PREFERENCES.log_prompt:
         formatted_response = "\n" + response
         logger.info(
             f"Response (ID: {request_id}): "

common/model.py

@@ -153,7 +153,7 @@ async def unload_embedding_model():
 def get_config_default(key: str, model_type: str = "model"):
     """Fetches a default value from model config if allowed by the user."""
 
-    default_keys = unwrap(config.model.get("use_as_default"), [])
+    default_keys = unwrap(config.model.use_as_default, [])
 
     # Add extra keys to defaults
     default_keys.append("embeddings_device")

common/networking.py

@@ -39,7 +39,7 @@ def handle_request_error(message: str, exc_info: bool = True):
     """Log a request error to the console."""
 
     trace = traceback.format_exc()
-    send_trace = unwrap(config.network.get("send_tracebacks"), False)
+    send_trace = config.network.send_tracebacks
 
     error_message = TabbyRequestErrorMessage(
         message=message, trace=trace if send_trace else None
@@ -134,7 +134,7 @@ def get_global_depends():
 
     depends = [Depends(add_request_id)]
 
-    if config.logging.get("requests"):
+    if config.logging.log_requests:
         depends.append(Depends(log_request))
 
     return depends

common/tabby_config.py

@@ -4,21 +4,11 @@ from loguru import logger
 from typing import Optional
 
 from common.utils import unwrap, merge_dicts
+from common.config_models import tabby_config_model
+import common.config_models
 
 
-class TabbyConfig:
-    network: dict = {}
-    logging: dict = {}
-    model: dict = {}
-    draft_model: dict = {}
-    lora: dict = {}
-    sampling: dict = {}
-    developer: dict = {}
-    embeddings: dict = {}
-
-    def __init__(self):
-        pass
-
+class TabbyConfig(tabby_config_model):
     def load_config(self, arguments: Optional[dict] = None):
         """load the global application config"""
 
@@ -30,14 +20,11 @@ class TabbyConfig:
 
         merged_config = merge_dicts(*configs)
 
-        self.network = unwrap(merged_config.get("network"), {})
-        self.logging = unwrap(merged_config.get("logging"), {})
-        self.model = unwrap(merged_config.get("model"), {})
-        self.draft_model = unwrap(merged_config.get("draft"), {})
-        self.lora = unwrap(merged_config.get("draft"), {})
-        self.sampling = unwrap(merged_config.get("sampling"), {})
-        self.developer = unwrap(merged_config.get("developer"), {})
-        self.embeddings = unwrap(merged_config.get("embeddings"), {})
+        for field in tabby_config_model.model_fields.keys():
+            value = unwrap(merged_config.get(field), {})
+            model = getattr(common.config_models, f"{field}_config_model")
+
+            setattr(self, field, model.parse_obj(value))
 
     def _from_file(self, config_path: pathlib.Path):
         """loads config from a given file path"""