From 33969b5845fec660d206ba9deba79de5500d4e6c Mon Sep 17 00:00:00 2001 From: Henk Date: Fri, 5 May 2023 17:23:01 +0200 Subject: [PATCH] Basic HF code execution support --- aiserver.py | 12 +++++++++--- koboldai_settings.py | 5 +++-- modeling/inference_models/hf_torch.py | 1 + 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/aiserver.py b/aiserver.py index 8d481b75..32823dd0 100644 --- a/aiserver.py +++ b/aiserver.py @@ -1345,7 +1345,8 @@ def general_startup(override_args=None): parser.add_argument("--summarizer_model", action='store', default="philschmid/bart-large-cnn-samsum", help="Huggingface model to use for summarization. Defaults to sshleifer/distilbart-cnn-12-6") parser.add_argument("--max_summary_length", action='store', default=75, help="Maximum size for summary to send to image generation") parser.add_argument("--multi_story", action='store_true', default=False, help="Allow multi-story mode (experimental)") - + parser.add_argument("--trust_remote_code", action='store_true', default=False, help="Allow Huggingface Models to Execute Code (Insecure!)") + parser.add_argument('-f', action='store', help="option for compatability with colab memory profiles") parser.add_argument('-v', '--verbosity', action='count', default=0, help="The default logging level is ERROR or higher. This value increases the amount of logging seen in your screen") parser.add_argument('-q', '--quiesce', action='count', default=0, help="The default logging level is ERROR or higher. This value decreases the amount of logging seen in your screen") @@ -1469,8 +1470,13 @@ def general_startup(override_args=None): allowed_ips = sorted(allowed_ips, key=lambda ip: int(''.join([i.zfill(3) for i in ip.split('.')]))) print(f"Allowed IPs: {allowed_ips}") - - + if args.trust_remote_code: + logger.warning("EXECUTION OF UNSAFE REMOTE CODE IS ENABLED!!!") + logger.warning("EXECUTION OF UNSAFE REMOTE CODE IS ENABLED!!!") + logger.warning("EXECUTION OF UNSAFE REMOTE CODE IS ENABLED!!!") + logger.warning("EXECUTION OF UNSAFE REMOTE CODE IS ENABLED!!!") + logger.warning("EXECUTION OF UNSAFE REMOTE CODE IS ENABLED!!!") + koboldai_vars.trust_remote_code = True if args.cpu: koboldai_vars.use_colab_tpu = False diff --git a/koboldai_settings.py b/koboldai_settings.py index dfccd4ef..d8416df2 100644 --- a/koboldai_settings.py +++ b/koboldai_settings.py @@ -1207,12 +1207,12 @@ class system_settings(settings): local_only_variables = ['lua_state', 'lua_logname', 'lua_koboldbridge', 'lua_kobold', 'lua_koboldcore', 'regex_sl', 'acregex_ai', 'acregex_ui', 'comregex_ai', 'comregex_ui', 'sp', '_horde_pid', 'inference_config', 'image_pipeline', - 'summarizer', 'summary_tokenizer', 'tts_model', 'rng_states', 'comregex_ai', 'comregex_ui'] + 'summarizer', 'summary_tokenizer', 'tts_model', 'rng_states', 'comregex_ai', 'comregex_ui', 'trust_remote_code'] no_save_variables = ['lua_state', 'lua_logname', 'lua_koboldbridge', 'lua_kobold', 'lua_koboldcore', 'sp', 'sp_length', '_horde_pid', 'horde_share', 'aibusy', 'serverstarted', 'inference_config', 'image_pipeline', 'summarizer', 'summary_tokenizer', 'use_colab_tpu', 'noai', 'disable_set_aibusy', 'cloudflare_link', 'tts_model', - 'generating_image', 'bit_8_available', 'host', 'hascuda', 'usegpu', 'rng_states', 'comregex_ai', 'comregex_ui', 'git_repository', 'git_branch'] + 'generating_image', 'bit_8_available', 'host', 'hascuda', 'usegpu', 'rng_states', 'comregex_ai', 'comregex_ui', 'git_repository', 'git_branch', 'trust_remote_code'] settings_name = "system" def __init__(self, socketio, koboldai_var): self._socketio = socketio @@ -1298,6 +1298,7 @@ class system_settings(settings): self.seen_messages = [] self.git_repository = "" self.git_branch = "" + self.trust_remote_code = False @dataclass diff --git a/modeling/inference_models/hf_torch.py b/modeling/inference_models/hf_torch.py index 49cdfc0f..ca1f1cdf 100644 --- a/modeling/inference_models/hf_torch.py +++ b/modeling/inference_models/hf_torch.py @@ -265,6 +265,7 @@ class HFTorchInferenceModel(HFInferenceModel): def _get_model(self, location: str, tf_kwargs: Dict): tf_kwargs["revision"] = utils.koboldai_vars.revision tf_kwargs["cache_dir"] = "cache" + tf_kwargs["trust_remote_code"] = utils.koboldai_vars.trust_remote_code # If we have model hints for legacy model, use them rather than fall back. try: