From 70810380c5be8c71749e57c68a9c643d2c940e50 Mon Sep 17 00:00:00 2001 From: rchan Date: Wed, 3 Jan 2024 14:55:31 +0000 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20convert=20types=20in=20setup=5Fl?= =?UTF-8?q?lm=20rather=20than=20in=20parser?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- reginald/models/setup_llm.py | 55 +++++++++++++++++++++++++----------- reginald/parser_utils.py | 27 +++++++++--------- reginald/utils.py | 8 ++++-- 3 files changed, 57 insertions(+), 33 deletions(-) diff --git a/reginald/models/setup_llm.py b/reginald/models/setup_llm.py index 6de372a5..f976300c 100644 --- a/reginald/models/setup_llm.py +++ b/reginald/models/setup_llm.py @@ -28,13 +28,13 @@ def setup_llm( data_dir: str | None = None, which_index: str | None = None, force_new_index: bool | str | None = None, - max_input_size: int | None = None, - k: int | None = None, - chunk_size: int | None = None, - chunk_overlap_ratio: float | None = None, - num_output: int | None = None, + max_input_size: int | str | None = None, + k: int | str | None = None, + chunk_size: int | str | None = None, + chunk_overlap_ratio: float | str | None = None, + num_output: int | str | None = None, is_path: bool | str | None = None, - n_gpu_layers: int | None = None, + n_gpu_layers: int | str | None = None, device: str | None = None, ) -> ResponseModel: """ @@ -65,31 +65,37 @@ def setup_llm( Whether to recreate the index vector store or not, by default None (uses False). If this is a string, it is converted to a boolean using `force_new_index.lower() == "true"`. - max_input_size : int | None, optional + max_input_size : int | str | None, optional Select the maximum input size for the model, by default None (uses 4096). Ignored if not using "llama-index-llama-cpp" or - "llama-index-hf" models - k : int | None, optional + "llama-index-hf" models, If this is a string, it is converted + to an integer + k : int | str | None, optional `similarity_top_k` to use in chat or query engine, - by default None (uses 3) - chunk_size : int | None, optional + by default None (uses 3). If this is a string, it is converted + to an integer + chunk_size : int | str | None, optional Maximum size of chunks to use, by default None. If None, this is computed as `ceil(max_input_size / k)`. - chunk_overlap_ratio : float, optional - Chunk overlap as a ratio of chunk size, by default None (uses 0.1) - num_output : int | None, optional - Number of outputs for the LLM, by default None (uses 512) + If this is a string, it is converted to an integer + chunk_overlap_ratio : float | str | None, optional + Chunk overlap as a ratio of chunk size, by default None (uses 0.1). + If this is a string, it is converted to a float + num_output : int | str | None, optional + Number of outputs for the LLM, by default None (uses 512). + If this is a string, it is converted to an integer is_path : bool | str | None, optional Whether or not model_name is used as a path to the model file, otherwise it should be the URL to download the model, by default None (uses False). If this is a string, it is converted to a boolean using `is_path.lower() == "true"`. Ignored if not using "llama-index-llama-cpp" model - n_gpu_layers : int | None, optional + n_gpu_layers : int | str | None, optional Select the number of GPU layers to use. If -1, all layers are offloaded to the GPU. If 0, no layers are offloaded to the GPU, by default None (uses 0). Ignored if not using - "llama-index-llama-cpp" model + "llama-index-llama-cpp" model. If this is a string, it is + converted to an integer device : str | None, optional Select which device to use, by default None (uses "auto"). Ignored if not using "llama-index-llama-cpp" or "llama-index-hf" models @@ -138,18 +144,31 @@ def setup_llm( # default for max_input_size if max_input_size is None: max_input_size = DEFAULT_ARGS["max_input_size"] + if isinstance(max_input_size, str): + max_input_size = int(max_input_size) # default for k if k is None: k = DEFAULT_ARGS["k"] + if isinstance(k, str): + k = int(k) + + # convert chunk_size if provided as str + # default is computed later using values of max_input_size and k + if isinstance(chunk_size, str): + chunk_size = int(chunk_size) # default for chunk_overlap_ratio if chunk_overlap_ratio is None: chunk_overlap_ratio = DEFAULT_ARGS["chunk_overlap_ratio"] + if isinstance(chunk_overlap_ratio, str): + chunk_overlap_ratio = float(chunk_overlap_ratio) # default for num_output if num_output is None: num_output = DEFAULT_ARGS["num_output"] + if isinstance(num_output, str): + num_output = int(num_output) # default for is_path if is_path is None: @@ -160,6 +179,8 @@ def setup_llm( # default for n_gpu_layers if n_gpu_layers is None: n_gpu_layers = DEFAULT_ARGS["n_gpu_layers"] + if isinstance(n_gpu_layers, str): + n_gpu_layers = int(n_gpu_layers) # default for device if device is None: diff --git a/reginald/parser_utils.py b/reginald/parser_utils.py index 8a44bdd0..4f4db0b7 100644 --- a/reginald/parser_utils.py +++ b/reginald/parser_utils.py @@ -2,6 +2,7 @@ import pathlib from reginald.models.models import MODELS +from reginald.models.setup_llm import DEFAULT_ARGS from reginald.utils import get_env_var @@ -32,7 +33,9 @@ def __init__(self, create_index_only: bool = False, *args, **kwargs): "-m", type=str, help=("Select which type of model to use " "Default is 'hello'."), - default=lambda: get_env_var("REGINALD_MODEL", secret_value=False), + default=lambda: get_env_var( + "REGINALD_MODEL", secret_value=False, default="hello" + ), choices=MODELS, ) self.add_argument( @@ -95,8 +98,8 @@ def __init__(self, create_index_only: bool = False, *args, **kwargs): "(ignored if not using llama-index-llama-cpp). " "Default is 0." ), - default=lambda: int( - get_env_var("LLAMA_INDEX_N_GPU_LAYERS", secret_value=False) + default=lambda: get_env_var( + "LLAMA_INDEX_N_GPU_LAYERS", secret_value=False ), ) self.add_argument( @@ -145,8 +148,8 @@ def __init__(self, create_index_only: bool = False, *args, **kwargs): "(ignored if not using llama-index). " "Default is 4096." ), - default=lambda: int( - get_env_var("LLAMA_INDEX_MAX_INPUT_SIZE", secret_value=False) + default=lambda: get_env_var( + "LLAMA_INDEX_MAX_INPUT_SIZE", secret_value=False ), ) self.add_argument( @@ -157,7 +160,7 @@ def __init__(self, create_index_only: bool = False, *args, **kwargs): "(ignored if not using llama-index). " "Default is 3." ), - default=lambda: int(get_env_var("LLAMA_INDEX_K", secret_value=False)), + default=lambda: get_env_var("LLAMA_INDEX_K", secret_value=False), ) self.add_argument( "--chunk-size", @@ -168,9 +171,7 @@ def __init__(self, create_index_only: bool = False, *args, **kwargs): "(ignored if not using llama-index). " "Default is computed by ceil(max_input_size / k)." ), - default=lambda: int( - get_env_var("LLAMA_INDEX_CHUNK_SIZE", secret_value=False) - ), + default=lambda: get_env_var("LLAMA_INDEX_CHUNK_SIZE", secret_value=False), ) self.add_argument( "--chunk-overlap-ratio", @@ -181,8 +182,8 @@ def __init__(self, create_index_only: bool = False, *args, **kwargs): "(ignored if not using llama-index). " "Default is 0.1." ), - default=lambda: float( - get_env_var("LLAMA_INDEX_CHUNK_OVERLAP_RATIO", secret_value=False) + default=lambda: get_env_var( + "LLAMA_INDEX_CHUNK_OVERLAP_RATIO", secret_value=False ), ) self.add_argument( @@ -194,9 +195,7 @@ def __init__(self, create_index_only: bool = False, *args, **kwargs): "(ignored if not using llama-index). " "Default is 512." ), - default=lambda: int( - get_env_var("LLAMA_INDEX_NUM_OUTPUT", secret_value=False) - ), + default=lambda: get_env_var("LLAMA_INDEX_NUM_OUTPUT", secret_value=False), ) diff --git a/reginald/utils.py b/reginald/utils.py index 18dffee7..0bd0498e 100644 --- a/reginald/utils.py +++ b/reginald/utils.py @@ -2,7 +2,9 @@ import os -def get_env_var(var: str, log: bool = True, secret_value: bool = True) -> str | None: +def get_env_var( + var: str, log: bool = True, secret_value: bool = True, default: str = None +) -> str | None: """ Get environment variable. Logs provided if log is True. @@ -16,6 +18,8 @@ def get_env_var(var: str, log: bool = True, secret_value: bool = True) -> str | Whether or not the value is a secret, by default True. If True, the value will not be logged. Ignored if log is False. + default : str, optional + Default value if environment variable is not found, by default None Returns ------- @@ -24,7 +28,7 @@ def get_env_var(var: str, log: bool = True, secret_value: bool = True) -> str | """ if log: logging.info(f"Trying to get environment variable '{var}'") - value = os.getenv(var) + value = os.getenv(var, default=default) if log: if value is not None: