From 8b299525fdba75c2967697b5a486aed004282897 Mon Sep 17 00:00:00 2001
From: vfbd <ve.forbryderne@gmail.com>
Date: Thu, 11 Aug 2022 14:57:13 -0400
Subject: [PATCH] sendtoapi now automatically detects tokenizer

---
 aiserver.py | 50 ++++++++++++++++++++++++++++----------------------
 1 file changed, 28 insertions(+), 22 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index 614262e1..f1be6b6f 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -369,6 +369,7 @@ class vars:
     use_colab_tpu = os.environ.get("COLAB_TPU_ADDR", "") != "" or os.environ.get("TPU_NAME", "") != ""  # Whether or not we're in a Colab TPU instance or Kaggle TPU instance and are going to use the TPU rather than the CPU
     revision    = None
     standalone = False
+    api_tokenizer_id = None
     disable_set_aibusy = False
     disable_input_formatting = False
     disable_output_formatting = False
@@ -1259,7 +1260,6 @@ def general_startup(override_args=None):
     parser.add_argument("--override_rename", action='store_true', help="Renaming stories from inside the browser is disabled if you are using --remote and enabled otherwise. Using this option will instead allow renaming stories if using --remote and prevent renaming stories otherwise.")
     parser.add_argument("--configname", help="Force a fixed configuration name to aid with config management.")
     parser.add_argument("--colab", action='store_true', help="Optimize for Google Colab.")
-    parser.add_argument("--tokenizer", type=str, help="When using the \"KoboldAI API\" backend option, this controls the tokenizer to use. This can be set to a Hugging Face model ID or the path to a folder under \"models\" in the KoboldAI folder.")
     parser.add_argument("--nobreakmodel", action='store_true', help="Disables Breakmodel support completely.")
     parser.add_argument("--unblock", action='store_true', default=False, help="Unblocks the KoboldAI port to be accessible from other machines without optimizing for remote play (It is recommended to use --host instead)")
     parser.add_argument("--quiet", action='store_true', default=False, help="If present will suppress any story related text from showing on the console")
@@ -2539,30 +2539,10 @@ def load_model(use_gpu=True, gpu_layers=None, disk_layers=None, initial_load=Fal
             }
 
         # If we're running Colab or OAI, we still need a tokenizer.
-        if(vars.model == "Colab"):
+        if(vars.model in ("Colab", "API")):
             from transformers import GPT2TokenizerFast
             tokenizer = GPT2TokenizerFast.from_pretrained("EleutherAI/gpt-neo-2.7B", revision=vars.revision, cache_dir="cache")
             loadsettings()
-        elif(vars.model == "API"):
-            tokenizer_id = getattr(args, "tokenizer", None)
-            if tokenizer_id is None:
-                tokenizer_id = "EleutherAI/gpt-neo-2.7B"
-            if(os.path.isdir(tokenizer_id)):
-                try:
-                    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, revision=vars.revision, cache_dir="cache")
-                except:
-                    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, revision=vars.revision, cache_dir="cache", use_fast=False)
-            elif(os.path.isdir("models/{}".format(args.tokenizer.replace('/', '_')))):
-                try:
-                    tokenizer = AutoTokenizer.from_pretrained("models/{}".format(tokenizer_id.replace('/', '_')), revision=vars.revision, cache_dir="cache")
-                except:
-                    tokenizer = AutoTokenizer.from_pretrained("models/{}".format(tokenizer_id.replace('/', '_')), revision=vars.revision, cache_dir="cache", use_fast=False)
-            else:
-                try:
-                    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, revision=vars.revision, cache_dir="cache")
-                except:
-                    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, revision=vars.revision, cache_dir="cache", use_fast=False)
-            loadsettings()
         elif(vars.model == "OAI"):
             from transformers import GPT2TokenizerFast
             tokenizer = GPT2TokenizerFast.from_pretrained("gpt2", revision=vars.revision, cache_dir="cache")
@@ -3912,6 +3892,32 @@ def actionsubmit(data, actionmode=0, force_submit=False, force_prompt_gen=False,
     while(True):
         set_aibusy(1)
 
+        if(vars.model == "API"):
+            global tokenizer
+            tokenizer_id = requests.get(
+                vars.colaburl[:-8] + "/api/v1/model",
+            ).json()["result"]
+            if tokenizer_id != vars.api_tokenizer_id:
+                try:
+                    if(os.path.isdir(tokenizer_id)):
+                        try:
+                            tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, revision=vars.revision, cache_dir="cache")
+                        except:
+                            tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, revision=vars.revision, cache_dir="cache", use_fast=False)
+                    elif(os.path.isdir("models/{}".format(tokenizer_id.replace('/', '_')))):
+                        try:
+                            tokenizer = AutoTokenizer.from_pretrained("models/{}".format(tokenizer_id.replace('/', '_')), revision=vars.revision, cache_dir="cache")
+                        except:
+                            tokenizer = AutoTokenizer.from_pretrained("models/{}".format(tokenizer_id.replace('/', '_')), revision=vars.revision, cache_dir="cache", use_fast=False)
+                    else:
+                        try:
+                            tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, revision=vars.revision, cache_dir="cache")
+                        except:
+                            tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, revision=vars.revision, cache_dir="cache", use_fast=False)
+                except:
+                    print(f"WARNING:  Unknown tokenizer {repr(tokenizer_id)}")
+                vars.api_tokenizer_id = tokenizer_id
+
         if(disable_recentrng):
             vars.recentrng = vars.recentrngm = None