Merge commit 'refs/pull/82/head' of https://github.com/ebolam/KoboldAI into UI2

aiserver.py

@@ -821,8 +821,10 @@ def loadmodelsettings():
     if("nobreakmodel" in js):
         koboldai_vars.nobreakmodel = js["nobreakmodel"]
     if("sampler_order" in js):
-        koboldai_vars.sampler_order = js["sampler_order"]
+        sampler_order = koboldai_vars.sampler_order
-        koboldai_vars.default_preset['sampler_order'] = js["sampler_order"]
+        if(len(sampler_order) < 7):
+            sampler_order = [6] + sampler_order
+        koboldai_vars.sampler_order = sampler_order
     if("temp" in js):
         koboldai_vars.temp       = js["temp"]
         koboldai_vars.default_preset['temp'] = js["temp"]
@@ -965,7 +967,10 @@ def processsettings(js):
     if("andepth" in js):
         koboldai_vars.andepth    = js["andepth"]
     if("sampler_order" in js):
-        koboldai_vars.sampler_order = js["sampler_order"]
+        sampler_order = koboldai_vars.sampler_order
+        if(len(sampler_order) < 7):
+            sampler_order = [6] + sampler_order
+        koboldai_vars.sampler_order = sampler_order
     if("temp" in js):
         koboldai_vars.temp       = js["temp"]
     if("top_p" in js):
@@ -1363,23 +1368,25 @@ def get_model_info(model, directory=""):
     
 
 def get_layer_count(model, directory=""):
-    if(model not in ["InferKit", "Colab", "OAI", "GooseAI" , "ReadOnly", "TPUMeshTransformerGPTJ"]):
+    if(model not in ["InferKit", "Colab", "API", "OAI", "GooseAI" , "ReadOnly", "TPUMeshTransformerGPTJ"]):
-        if(koboldai_vars.model == "GPT2Custom"):
+        if(model == "GPT2Custom"):
-            model_config = open(directory + "/config.json", "r")
+            with open(os.path.join(directory, "config.json"), "r") as f:
+                model_config = json.load(f)
         # Get the model_type from the config or assume a model type if it isn't present
         else:
+            if(directory):
+                model = directory
             from transformers import AutoConfig
-            if directory == "":
+            if(os.path.isdir(model.replace('/', '_'))):
-                model_config = AutoConfig.from_pretrained(model, revision=koboldai_vars.revision, cache_dir="cache")
+                model_config = AutoConfig.from_pretrained(model.replace('/', '_'), revision=koboldai_vars.revision, cache_dir="cache")
+            elif(os.path.isdir("models/{}".format(model.replace('/', '_')))):
+                model_config = AutoConfig.from_pretrained("models/{}".format(model.replace('/', '_')), revision=koboldai_vars.revision, cache_dir="cache")
             elif(os.path.isdir(directory)):
                 model_config = AutoConfig.from_pretrained(directory, revision=koboldai_vars.revision, cache_dir="cache")
             elif(os.path.isdir(koboldai_vars.custmodpth.replace('/', '_'))):
                 model_config = AutoConfig.from_pretrained(koboldai_vars.custmodpth.replace('/', '_'), revision=koboldai_vars.revision, cache_dir="cache")
             else:
-                model_config = AutoConfig.from_pretrained(koboldai_vars.custmodpth, revision=koboldai_vars.revision, cache_dir="cache")
+                model_config = AutoConfig.from_pretrained(model, revision=koboldai_vars.revision, cache_dir="cache")
-        
-        
-        
         return utils.num_layers(model_config)
     else:
         return None
@@ -1623,8 +1630,6 @@ def patch_transformers():
     dynamic_processor_wrap(TailFreeLogitsWarper, "tfs", "tfs", cond=lambda x: x < 1.0)
     dynamic_processor_wrap(TypicalLogitsWarper, "typical", "typical", cond=lambda x: x < 1.0)
     dynamic_processor_wrap(TemperatureLogitsWarper, "temperature", "temp", cond=lambda x: x != 1.0)
-    RepetitionPenaltyLogitsProcessor.__init__ = AdvancedRepetitionPenaltyLogitsProcessor.__init__
-    RepetitionPenaltyLogitsProcessor.__call__ = AdvancedRepetitionPenaltyLogitsProcessor.__call__
 
     class PhraseBiasLogitsProcessor(LogitsProcessor):
         def __init__(self):
@@ -1768,9 +1773,13 @@ def patch_transformers():
             self.__warper_list.append(TailFreeLogitsWarper(tfs=0.5, min_tokens_to_keep=1 + (beams > 1)))
             self.__warper_list.append(TypicalLogitsWarper(typical=0.5, min_tokens_to_keep=1 + (beams > 1)))
             self.__warper_list.append(TemperatureLogitsWarper(temperature=0.5))
+            self.__warper_list.append(AdvancedRepetitionPenaltyLogitsProcessor())
 
         def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, *args, **kwargs):
-            for k in koboldai_vars.sampler_order:
+            sampler_order = koboldai_vars.sampler_order[:]
+            if len(sampler_order) < 7:  # Add repetition penalty at beginning if it's not present
+                sampler_order = [6] + sampler_order
+            for k in sampler_order:
                 scores = self.__warper_list[k](input_ids, scores, *args, **kwargs)
             return scores
 
@@ -2525,6 +2534,9 @@ def load_model(use_gpu=True, gpu_layers=None, disk_layers=None, initial_load=Fal
             koboldai_vars.compiling = False
         
         def tpumtjgenerate_settings_callback() -> dict:
+            sampler_order = vars.sampler_order[:]
+            if len(sampler_order) < 7:  # Add repetition penalty at beginning if it's not present
+                sampler_order = [6] + sampler_order
             return {
                 "sampler_order": koboldai_vars.sampler_order,
                 "top_p": float(koboldai_vars.top_p),
@@ -3685,12 +3697,16 @@ def get_message(msg):
         sendUSStatItems()
     elif(msg['cmd'] == 'samplers'):
         sampler_order = msg["data"]
+        sampler_order_min_length = 6
+        sampler_order_max_length = 7
         if(not isinstance(sampler_order, list)):
             raise ValueError(f"Sampler order must be a list, but got a {type(sampler_order)}")
-        if(len(sampler_order) != len(koboldai_vars.sampler_order)):
+        if(not (sampler_order_min_length <= len(sampler_order) <= sampler_order_max_length)):
-            raise ValueError(f"Sampler order must be a list of length {len(koboldai_vars.sampler_order)}, but got a list of length {len(sampler_order)}")
+            raise ValueError(f"Sampler order must be a list of length greater than or equal to {sampler_order_min_length} and less than or equal to {sampler_order_max_length}, but got a list of length {len(sampler_order)}")
         if(not all(isinstance(e, int) for e in sampler_order)):
             raise ValueError(f"Sampler order must be a list of ints, but got a list with at least one non-int element")
+        if(min(sampler_order) != 0 or max(sampler_order) != len(sampler_order) - 1 or len(set(sampler_order)) != len(sampler_order)):
+            raise ValueError(f"Sampler order list of length {len(sampler_order)} must be a permutation of the first {len(sampler_order)} nonnegative integers")
         koboldai_vars.sampler_order = sampler_order
         settingschanged()
     elif(msg['cmd'] == 'list_model'):
@@ -3701,8 +3717,8 @@ def get_message(msg):
         changed = True
         if not utils.HAS_ACCELERATE:
             msg['disk_layers'] = "0"
-        if os.path.exists("settings/" + koboldai_vars.model.replace('/', '_') + ".breakmodel"):
+        if os.path.exists("settings/" + koboldai_vars.model_selected.replace('/', '_') + ".breakmodel"):
-            with open("settings/" + koboldai_vars.model.replace('/', '_') + ".breakmodel", "r") as file:
+            with open("settings/" + koboldai_vars.model_selected.replace('/', '_') + ".breakmodel", "r") as file:
                 data = file.read().split('\n')[:2]
                 if len(data) < 2:
                     data.append("0")
@@ -3710,14 +3726,15 @@ def get_message(msg):
                 if gpu_layers == msg['gpu_layers'] and disk_layers == msg['disk_layers']:
                     changed = False
         if changed:
-            if koboldai_vars.model in ["NeoCustom", "GPT2Custom"]:
+            if koboldai_vars.model_selected in ["NeoCustom", "GPT2Custom"]:
                 filename = "settings/{}.breakmodel".format(os.path.basename(os.path.normpath(koboldai_vars.custmodpth)))
             else:
-                filename = "settings/{}.breakmodel".format(koboldai_vars.model.replace('/', '_'))
+                filename = "settings/{}.breakmodel".format(koboldai_vars.model_selected.replace('/', '_'))
             f = open(filename, "w")
             f.write(str(msg['gpu_layers']) + '\n' + str(msg['disk_layers']))
             f.close()
         koboldai_vars.colaburl = msg['url'] + "/request"
+        vars.model = vars.model_selected
         load_model(use_gpu=msg['use_gpu'], gpu_layers=msg['gpu_layers'], disk_layers=msg['disk_layers'], online_model=msg['online_model'])
     elif(msg['cmd'] == 'show_model'):
         print("Model Name: {}".format(getmodelname()))
@@ -3742,18 +3759,18 @@ def get_message(msg):
         elif msg['data'] in ('NeoCustom', 'GPT2Custom') and 'path_modelname' in msg:
             #Here the user entered custom text in the text box. This could be either a model name or a path.
             if check_if_dir_is_model(msg['path_modelname']):
-                koboldai_vars.model = msg['data']
+                koboldai_vars.model_selected = msg['data']
                 koboldai_vars.custmodpth = msg['path_modelname']
                 get_model_info(msg['data'], directory=msg['path'])
             else:
-                koboldai_vars.model = msg['path_modelname']
+                koboldai_vars.model_selected = msg['path_modelname']
                 try:
-                    get_model_info(koboldai_vars.model)
+                    get_model_info(koboldai_vars.model_selected)
                 except:
                     emit('from_server', {'cmd': 'errmsg', 'data': "The model entered doesn't exist."}, room="UI_1")
         elif msg['data'] in ('NeoCustom', 'GPT2Custom'):
             if check_if_dir_is_model(msg['path']):
-                koboldai_vars.model = msg['data']
+                koboldai_vars.model_selected = msg['data']
                 koboldai_vars.custmodpth = msg['path']
                 get_model_info(msg['data'], directory=msg['path'])
             else:
@@ -3762,12 +3779,12 @@ def get_message(msg):
                 else:
                     sendModelSelection(menu=msg['data'], folder=msg['path'])
         else:
-            koboldai_vars.model = msg['data']
+            koboldai_vars.model_selected = msg['data']
             if 'path' in msg:
                 koboldai_vars.custmodpth = msg['path']
                 get_model_info(msg['data'], directory=msg['path'])
             else:
-                get_model_info(koboldai_vars.model)
+                get_model_info(koboldai_vars.model_selected)
     elif(msg['cmd'] == 'delete_model'):
         if "{}/models".format(os.getcwd()) in os.path.abspath(msg['data']) or "{}\\models".format(os.getcwd()) in os.path.abspath(msg['data']):
             if check_if_dir_is_model(msg['data']):
@@ -3873,7 +3890,6 @@ def get_message(msg):
             emit(
                 'from_server',
                 {'cmd': 'showfieldbudget', 'data': {"length": None, "max": None, "field": field}},
-                broadcast=True
             )
             return
 
@@ -4612,7 +4628,7 @@ def _generate(txt, minimum, maximum, found_entries):
                 gen_in, 
                 do_sample=True, 
                 max_length=int(2e9),
-                repetition_penalty=1.1,
+                repetition_penalty=1.0,
                 bad_words_ids=koboldai_vars.badwordsids,
                 use_cache=True,
                 num_return_sequences=numseqs

colab/TPU.ipynb

@@ -66,7 +66,7 @@
         "#@title <b><-- Select your model below and then click this to start KoboldAI</b>\n",
         "#@markdown You can find a description of the models below along with instructions on how to start KoboldAI.\n",
         "\n",
-        "Model = \"Nerys 13B V2\" #@param [\"Nerys 13B V2\", \"Janeway 13B\", \"Shinen 13B\", \"Skein 6B\", \"Janeway 6B\", \"Adventure 6B\", \"Shinen 6B\", \"Lit 6B\", \"NeoX 20B\", \"OPT 13B\", \"Fairseq Dense 13B\", \"GPT-J-6B\"] {allow-input: true}\n",
+        "Model = \"Nerys 13B V2\" #@param [\"Nerys 13B V2\", \"Janeway 13B\", \"Shinen 13B\", \"Skein 20B\", \"Skein 6B\", \"Janeway 6B\", \"Adventure 6B\", \"Shinen 6B\", \"Lit 6B\", \"NeoX 20B\", \"OPT 13B\", \"Fairseq Dense 13B\", \"GPT-J-6B\"] {allow-input: true}\n",
         "Version = \"Official\" #@param [\"Official\", \"United\"] {allow-input: true}\n",
         "Provider = \"Cloudflare\" #@param [\"Localtunnel\", \"Cloudflare\"]\n",
         "\n",
@@ -93,6 +93,10 @@
         "  Model = \"KoboldAI/fairseq-dense-13B-Shinen\"\n",
         "  path = \"\"\n",
         "  download = \"\"\n",
+        "elif Model == \"Skein 20B\":\n",
+        "  Model = \"KoboldAI/GPT-NeoX-20B-Skein\"\n",
+        "  path = \"\"\n",
+        "  download = \"\"\n",
         "elif Model == \"NeoX 20B\":\n",
         "  Model = \"EleutherAI/gpt-neox-20b\"\n",
         "  path = \"\"\n",
@@ -128,7 +132,7 @@
         "elif Model == \"GPT-J-6B\":\n",
         "  Model = \"EleutherAI/gpt-j-6B\"\n",
         "  path = \"\"\n",
-        "  download = \"\"\n",        
+        "  download = \"\"\n",
         "else:\n",
         "  path = \"\"\n",
         "  download = \"\"\n",
@@ -225,4 +229,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
+}

gensettings.py

@@ -495,6 +495,17 @@ gensettingstf = [
     "classname": "story",
     "name": "andepth"
  	},
+    {
+	"uitype": "toggle",
+	"unit": "bool",
+	"label": "Show Field Budget",
+	"id": "setshowbudget",
+	"min": 0,
+	"max": 1,
+	"step": 1,
+	"default": 0,
+  "tooltip": "Shows token usage when typing in relevant text boxes. <b>May lag slower devices.</b>"
+	},
 ]
 
 gensettingsik =[{

static/application.js

@@ -241,8 +241,27 @@ function addSetting(ob) {
 			if(ob.id == "setadventure"){
 				setadventure($(this).prop('checked'));
 			}
+
 		});
 	}
+
+	if (ob.id === "setshowbudget") {
+		$("#setshowbudget").on("change", function () {
+			for (const el of document.getElementsByClassName("input-token-usage")) {
+				if (this.checked) {
+					el.classList.remove("hidden");
+				} else {
+					el.classList.add("hidden");
+				}
+			}
+		});
+
+		if (!$("#setshowbudget")[0].checked) {
+			for (const el of document.getElementsByClassName("input-token-usage")) {
+				el.classList.add("hidden");
+			}
+		}
+	}
 }
 
 function refreshTitle() {
@@ -1287,12 +1306,13 @@ function buildSamplerList(samplers) {
 		"Tail-free Sampling",
 		"Typical Sampling",
 		"Temperature",
+		"Repetition Penalty",
 	]
 	for(i=0; i<samplers.length; i++) {
 		samplerslist.append("<div class=\"flex\">\
 			<div class=\"samplerslistitem flex-row-container\" sid=\""+samplers[i]+"\">\
 				<div class=\"flex-row\">\
-					<div>"+samplers_lookup_table[samplers[i]]+"</div>\
+					<div>"+(samplers[i] < samplers_lookup_table.length ? samplers_lookup_table[samplers[i]] : "Unknown sampler #" + samplers[i])+"</div>\
 				</div>\
 			</div>\
 		</div>");
@@ -2165,6 +2185,9 @@ function interpolateRGB(color0, color1, t) {
 }
 
 function updateInputBudget(inputElement) {
+	let budgetElement = document.getElementById("setshowbudget");
+	if (budgetElement && !budgetElement.checked) return;
+
 	let data = {"unencoded": inputElement.value, "field": inputElement.id};
 
 	if (inputElement.id === "anoteinput") {
@@ -2182,7 +2205,6 @@ function registerTokenCounters() {
 
 		let span = document.createElement("span");
 		span.classList.add("input-token-usage");
-		span.innerText = "?/? Tokens";
 		el.appendChild(span);
 
 		let inputElement = el.querySelector("input, textarea");
@@ -2446,10 +2468,6 @@ $(document).ready(function(){
 		} else if(msg.cmd == "updatechunk") {
 			hideMessage();
 			game_text.attr('contenteditable', allowedit);
-			if (typeof submit_start !== 'undefined') {
-				$("#runtime")[0].innerHTML = `Generation time: ${Math.round((Date.now() - submit_start)/1000)} sec`;
-				delete submit_start;
-			}
 			var index = msg.data.index;
 			var html = msg.data.html;
 			var existingChunk = game_text.children('#n' + index);
@@ -2963,6 +2981,7 @@ $(document).ready(function(){
 			$("#showmodelnamecontainer").removeClass("hidden");
 		} else if(msg.cmd == 'hide_model_name') {
 			$("#showmodelnamecontainer").addClass("hidden");
+			$(window).off('beforeunload');
 			location.reload();
 			//console.log("Closing window");
 		} else if(msg.cmd == 'model_load_status') {

static/custom.css

@@ -473,7 +473,7 @@ body.connected #popupfooter, #popupfooter.always-available {
 }
 
 #samplerslist {
-	height: 300px;
+	height: 310px;
 	overflow-y: scroll;
 	overflow-wrap: anywhere;
 }

static/favicon.js

@@ -2,6 +2,7 @@
 var fav_icon2 = "data:image/x-icon;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAMAAAAoLQ9TAAAABGdBTUEAALGPC/xhBQAAACBjSFJNAAB6JgAAgIQAAPoAAACA6AAAdTAAAOpgAAA6mAAAF3CculE8AAAB+1BMVEUAAAAAAAAAAAAAAAAAAQAAAAAAAQAAAAAAAAASFhBBWD4iUyoFEwgFEwguUTM+VDoMFAwAAAA+elIudz8AAAAAAAA0MigyLyQAAAAbLh1LdElSbUoVMBkAAABAZ0M2fkUAAAABAQFMiGQraDkAAQANFxEGFQkLFg8EEAYAAAAsZDonZjUAAABCgVVAnFYrSjhEjFpFi1sdRScAAAAjOi8VMxx1dGOFgGYAAABOTEabmIdlYlQaGhgaGhddXFauqY5JRjoAAAAAAAABAQFGeExIl1lX0XRW0XRHi1RFe02vv5W31KFd1Hpc1Hpe1HvO1KvDvJlqZ1plYVOmoIVt1IFl1H7AuZp1cV9jX1AmSCw3Nzg7NmA1MTJuz4Bm1H5MST9HPl9BQEMgNiNXgWKiobFgXICDd5dfw3RZVnJiV3zGv9Bqf29Oj2G/v8hTTpGhl8dbxHVd0npiYoxhWJvIxtlcimZFn1lRclg9SkZNblZBeEpDbEZCa0ZBc0hLY1BAS1BdaV87j01Vx3FWynJSrGZOhlVasGtas2xatm1at21WnWJQm15WyXJQvmlavnBZrGlEYEJWe1RBWz9Um2BavXBgxn9XhllGY0RLaklXiFlTwG5OpmVSfFNMbUpGZEVLa0lShldEhVCChHiKiHvWz6/Kw6WWlZGAfmj///8kr0X+AAAARHRSTlMAASFrcAhxIjLb/vWvsPb+20b4+DFFyMkz2vf43CP9/m5y9vZysLGvsQn19mz+/tz4+NxHycr3+Ejb/vaxsPX+3TRtcBrzrrgAAAABYktHRKhQCDaSAAAAB3RJTUUH5gYJFyQy3tftxgAAAQBJREFUGNNjYGBgYGRiZmFlZWNmZ2SAAA5OLm4eXj5+AQ6ogKCQi6ubu4ensCCIxygiKubl7ePr6+cfIC4owcjAJCkVGBQc4usbGhYeIS0jy8AsFxkVHRPr6xsXn5CYJK/AoKiUnJKalg5UkZGZla2swsCqmpObl1/g61tYVFxSqsbKwKpeVl5RWVVdU1tX39CoocnAotXU3NLa1t7R2dXd06utwqCj6+vb1z9h4sRJk6f4+uopMLDrG0z1nTZ94sQZM31nGRrJMjBKGJvMnjN3wrz5CxaaCnKAvSNqtmjxkqXLlptbQP0iYmllbWNrZ+/gCBVgZHdS1GR1VpAFqQcApI0/jqlZOvEAAAAldEVYdGRhdGU6Y3JlYXRlADIwMjItMDYtMDlUMjM6MzY6NTArMDA6MDDi0xr+AAAAJXRFWHRkYXRlOm1vZGlmeQAyMDIyLTA2LTA5VDIzOjM2OjUwKzAwOjAwk46iQgAAAABJRU5ErkJggg==";
 var fav_icon1 = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAMAAAAoLQ9TAAAABGdBTUEAALGPC/xhBQAAACBjSFJNAAB6JgAAgIQAAPoAAACA6AAAdTAAAOpgAAA6mAAAF3CculE8AAAB+FBMVEUAAAAAAAAAAAAAAAAAAAEAAAAAAQEAAAAAAAAUFRlLVGYrSWgHEBoHEBk3S19HUGMOExkAAABOcos7apIAAAAAAAA2Ly01KyoAAAAgKzdVaX9bZHIaKzwAAABKYHhDcZgAAAABAQFfgJY2XX0AAQEQFhoIEhwOFRgGDRUAAAAAAQE3W3cyWnwAAABSeJJRjLs1R1FVgaFWgJ4lPlMAAAAsOD4aLj55bm2Md3QAAABPSkmfko9pXlsbGRkbGRlfWlm1oJxMQkAAAAAAAAABAQFTb4tYibFtvPpWgKNScpC6s7nExtNzwPp1wPnZx8jMsKtuZGFoXVutmJODwfJ7wfbHr6p5a2hnW1gtQlI4ODk7N2A2LzWDvet8wPZPRkRHPl9CQUQlMTthe4+ko7RhXYGEeJhzsuJaVXRjWHzIwtNwfYddhqLCwcpTTpGimMhvsuVzv/djYpBgWJvLydxlgptVirdZbX1ASFZUaXtOb4xOZX1OZHxNa4ZRX21DSV5gaG9Je6lqsepstO1knclcfJxtoc5tpNFuptVup9ZnkbdgjrVss+xjpuBvrd9snspOW29jdI5LVmlkj7Vvrd54t+RlfptQXXJWZHtlf51oruNgmMFfdJBYZn1RXnRWZXthfZxSeZiGgYGOhYLdxb/RubWZlpWFd3T////2kwjgAAAARXRSTlMAASFrcAhxIjLb/vWvsPb+20b4+DFFyMkz2vf43CP9/m5y9vZysLGvsQlw9fZs/v7c+PjcR8nK9/hI2/72sbD1/t00bXBAFktiAAAAAWJLR0SnwLcrAwAAAAd0SU1FB+YGCRchHQhxJNoAAAD/SURBVBjTY2BgYGBkYmZhZWVjZmdkgAAOTi5uHl4+fgEOqICgkKubu7uHp7AgiMcoIirm5e3j4+Pr5y8uKMHIwCQpFRAYFOzjExIaFi4tI8vALBcRGRUd4+MTGxefkCivwKColJSckpoGVJGekZmlrMLAqpqdk5uX7+NTUFhUXKLGysCqXlpWXlFZVV1TW1ffoKHJoKXd2NTc0trW3tHZ1d2jo8Kgq+fj09vXP2HCxEmTfXz0FRjYDQyn+EydNmHC9Bk+M42MZRkYJUxMZ82e0z933vwFZoIcYO+Imi9ctHjJ0mUWllC/iFhZ29ja2Ts4OkEFGNmdFTVZXRRkQeoBhkE/Yj5NSZ4AAAAldEVYdGRhdGU6Y3JlYXRlADIwMjItMDYtMDlUMjM6MzM6MjgrMDA6MDA90JbEAAAAJXRFWHRkYXRlOm1vZGlmeQAyMDIyLTA2LTA5VDIzOjMzOjI4KzAwOjAwTI0ueAAAAABJRU5ErkJggg==";
 var fav_icon = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAMAAAAoLQ9TAAAABGdBTUEAALGPC/xhBQAAACBjSFJNAAB6JgAAgIQAAPoAAACA6AAAdTAAAOpgAAA6mAAAF3CculE8AAAB8lBMVEUAAAAAAAAAAAAAAAABAAAAAAABAAAAAAAAAAAdEBB0Pz5rKCgaBwcZBwdkMzJxPDocDAwAAACLTU6SOzsAAAAAAAA9Mic/LyEAAAA6HByQUUaIVEY+GBgAAACAQkKaQUIAAAABAQGWXl9+NjYBAAAaEBAcCAgZDQ0WBQUAAAB3Nzd9MjIAAACTUVK7UVJRNTWhVVaeVldTJSUAAAA+LC0+GhuGcmCgf2EAAABUTESrl4NzYlEdGhcdGhdiXFbIqIhWRjcAAAAAAAABAQGUSkq1VVX6bW6oUVGXS0vmro7+uJn6c3T6dXX/yqPnu5F3aFhxYVG/oH/7gHv6enjeuJOEcFtzX01VLCs4ODk7NmA5MTH1gHr6e3hWSTxHPl9CQUQ/JCKPYGGko7RhXYGEeJjmcW9cVnFjWH3IwtOHb3CjXV3CwcpTTpGimMjlb3D4c3RmYI1gWJvLydybZWW+T0x+V1hRP0Z7U1WTSEiHRUWGRUSORkZuTlBRQVBwX2CvRkXtaGjvamrNYWKmU1PVZ2fXaGjbaWncaWnAX1+7W1vkYF/ja2zRZWV9QkGeVFN2Pz69XV3ia2zkeHmpWFd/REOJSUirWVjjaGjBYGCeUlKMSkl8QkGBRUSoVlWeUE2QgXeWiHr1zqjmw5+bl5KVe2T///8NZLRGAAAARHRSTlMAASFrcAhxIjLb/vWvsPb+20b4+DFFyMkz2vf43CP9/m5y9vZysLGvsQn19mz+/tz4+NxHycr3+Ejb/vaxsPX+3TRtcBrzrrgAAAABYktHRKUuuUovAAAAB3RJTUUH5gYJFzsfVlK/LQAAAP9JREFUGNNjYGBgYGRiZmFlZWNmZ2SAAA5OLm4eXj5+AQ6ogKCQi6ubm7uHsCCIxygiKubp5e3t7ePrJy4owcjAJCnlHxAY5O0dHBIaJi0jy8AsFx4RGRXt7R0TGxefIK/AoKiUmJSckgpUkZaekamswsCqmpWdk5vn7Z1fUFhUrMbKwKpeUlpWXlFZVV1TW1evocnAotXQ2NTc0trW3tHZ2KWtwqCj6+3d3dPb19c/YaK3t54CA7u+wSTvyVP6+qZO855uaCTLwChhbDJj5qzZc6bOnWcqyAH2jqjZ/AULFy1eYm4B9YuIpZW1ja2dvYMjVICR3UlRk9VZQRakHgAlRz6K4dvoSgAAACV0RVh0ZGF0ZTpjcmVhdGUAMjAyMi0wNi0wOVQyMzo1OTozMSswMDowMJt1iQMAAAAldEVYdGRhdGU6bW9kaWZ5ADIwMjItMDYtMDlUMjM6NTk6MzErMDA6MDDqKDG/AAAAAElFTkSuQmCC"
+var submit_start;
 
 var favicon = {
 
@@ -53,11 +54,16 @@ var favicon = {
 	start_swap: function() {
 		this.run = true;
 		this.auto_swap();
+		submit_start = Date.now();
 	},
 	
 	stop_swap: function() {
 		this.run = false;
 		this.change(fav_icon);
+		if (typeof submit_start !== 'undefined') {
+			$("#runtime")[0].innerHTML = `Execution time: ${Math.round((Date.now() - submit_start)/1000)} sec`;
+			delete submit_start;
+		}
 	},
 	
 	docHead:document.getElementsByTagName("head")[0]

tpu_mtj_backend.py

@@ -176,7 +176,7 @@ def apply_repetition_penalty_dynamic(logits, tokens, repetition_penalty, generat
     logits[tokens] = penalty_logits
     return logits
 
-def kobold_sample_dynamic(key, logits, sampler_order: Optional[np.ndarray] = None, top_p=0.9, temp=0.5, top_k=0, tfs=1.0, typical=1.0, top_a=0.0):
+def kobold_sample_dynamic(key, logits, rpargs, sampler_order: Optional[np.ndarray] = None, top_p=0.9, temp=0.5, top_k=0, tfs=1.0, typical=1.0, top_a=0.0):
     '''
     This gets called by generate_loop_fn to apply a series of 6 filters
     to the logits (top-k, then top-a, then top-p, then TFS, then typical, then temperature)
@@ -312,6 +312,7 @@ def kobold_sample_dynamic(key, logits, sampler_order: Optional[np.ndarray] = Non
         if k == 3 and tfs < 1.0: logits = tail_free_filter(logits)
         if k == 4 and typical < 1.0: logits = typical_filter(logits)
         if k == 5 and temp != 1.0: logits = temp_filter(logits)
+        if k == 6 and rpargs[1] != 1.0: logits = apply_repetition_penalty_dynamic(logits, *rpargs)
     # Finally, pick one token using the softmax thingy again (it gives
     # an array whose elements sum to 1 so it can be used nicely as a
     # probability distribution)
@@ -362,7 +363,7 @@ def apply_repetition_penalty_static(logits, tokens, repetition_penalty, generate
     # positions in the logits array
     return logits.at[tokens].set(penalty_logits)
 
-def kobold_sample_static(key, logits, sampler_order: Optional[np.ndarray] = None, top_p=0.9, temp=0.5, top_k=0, tfs=1.0, typical=1.0, top_a=0.0):
+def kobold_sample_static(key, logits, rpargs, sampler_order: Optional[np.ndarray] = None, top_p=0.9, temp=0.5, top_k=0, tfs=1.0, typical=1.0, top_a=0.0):
     '''
     This gets called by generate_loop_fn to apply a series of 6 filters
     to the logits (top-k, then top-a, then top-p, then TFS, then typical, then temperature)
@@ -497,6 +498,7 @@ def kobold_sample_static(key, logits, sampler_order: Optional[np.ndarray] = None
         logits = jax.lax.cond(jnp.logical_and(k == 3, tfs < 1.0), tail_free_filter, lambda x: x, logits)
         logits = jax.lax.cond(jnp.logical_and(k == 4, typical < 1.0), typical_filter, lambda x: x, logits)
         logits = jax.lax.cond(jnp.logical_and(k == 5, temp != 1.0), temp_filter, lambda x: x, logits)
+        logits = jax.lax.cond(jnp.logical_and(k == 6, rpargs[1] != 1.0), lambda x: apply_repetition_penalty_static(*x), lambda x: x[0], (logits, *rpargs))
     # Finally, pick one token using the softmax thingy again (it gives
     # an array whose elements sum to 1 so it can be used nicely as a
     # probability distribution)
@@ -513,17 +515,6 @@ def sample_func(data, key, numseqs_aux, badwords, repetition_penalty, generated_
         # Get the pseudo-random number generator key that will
         # be used by kobold_sample_dynamic to randomly pick a token
         sample_key, new_key = jax.random.split(sample_key, num=2)
-        # Apply repetition penalty to all tokens that are
-        # currently inside the "generated" array
-        logits = apply_repetition_penalty_dynamic(
-            logits,
-            generated,
-            repetition_penalty,
-            generated_index, 
-            gen_length,
-            rpslope,
-            rprange,
-        )
         # Remove any tokens in the badwords list by setting
         # their logits to negative infinity which effectively
         # makes their probabilities of being chosen zero
@@ -535,6 +526,14 @@ def sample_func(data, key, numseqs_aux, badwords, repetition_penalty, generated_
         next_token = kobold_sample_dynamic(
             sample_key,
             logits,
+            (
+                generated,
+                repetition_penalty,
+                generated_index, 
+                gen_length,
+                rpslope,
+                rprange,
+            ),
             **sampler_options,
         )
         # Remember what token was picked
@@ -606,18 +605,6 @@ class PenalizingCausalTransformer(CausalTransformer):
                     assert logits.shape == (1, config["n_vocab"])
                     # Flatten it into a 1D array to make it easier to use
                     logits = logits[0]
-                    # Apply repetition penalty to all tokens that are
-                    # currently inside the "generated" array
-                    if repetition_penalty is not None:
-                        logits = apply_repetition_penalty_static(
-                            logits,
-                            generated,
-                            repetition_penalty,
-                            generated_index,
-                            gen_length,
-                            rpslope,
-                            rprange,
-                        )
                     # Remove any tokens in the badwords list by setting
                     # their logits to negative infinity which effectively
                     # makes their probabilities of being chosen zero
@@ -629,6 +616,14 @@ class PenalizingCausalTransformer(CausalTransformer):
                     next_token = kobold_sample_static(
                         sample_key,
                         logits,
+                        (
+                            generated,
+                            repetition_penalty,
+                            generated_index,
+                            gen_length,
+                            rpslope,
+                            rprange,
+                        ),
                         **sampler_options,
                     )
                     # Remember what token was picked
@@ -863,6 +858,9 @@ def infer_static(
     maps.thread_resources.env = thread_resources_env
     if sampler_order is None:
         sampler_order = utils.default_sampler_order.copy()
+    sampler_order = sampler_order[:]
+    if len(sampler_order) < 7:  # Add repetition penalty at beginning if it's not present
+        sampler_order = [6] + sampler_order
     sampler_order = np.uint32(sampler_order)
     total_batch = 1
     tokens = context

utils.py

@@ -33,7 +33,7 @@ layers_module_names: Optional[List[str]] = None
 module_names: Optional[List[str]] = None
 named_buffers: Optional[List[tuple]] = None
 
-default_sampler_order = [0, 1, 2, 3, 4, 5]
+default_sampler_order = [6, 0, 1, 2, 3, 4, 5]
 
 #==================================================================#
 # Decorator to prevent a function's actions from being run until
@@ -167,7 +167,7 @@ def decodenewlines(txt):
 #  Returns number of layers given an HF model config
 #==================================================================#
 def num_layers(config):
-    return config.num_layers if hasattr(config, "num_layers") else config.n_layer if hasattr(config, "n_layer") else config.num_hidden_layers if hasattr(config, 'num_hidden_layers') else None
+    return config["n_layer"] if isinstance(config, dict) else config.num_layers if hasattr(config, "num_layers") else config.n_layer if hasattr(config, "n_layer") else config.num_hidden_layers if hasattr(config, 'num_hidden_layers') else None
 
 #==================================================================#
 #  Downloads huggingface checkpoints using aria2c if possible
@@ -177,6 +177,7 @@ class Send_to_socketio(object):
     def write(self, bar):
         time.sleep(0.01)
         try:
+            print(bar)
             emit('from_server', {'cmd': 'model_load_status', 'data': bar.replace(" ", " ")}, broadcast=True)
         except:
             pass

warpers.py

@@ -28,10 +28,10 @@ SOFTWARE.
 '''
 
 import torch
-from transformers import LogitsWarper, LogitsProcessor
+from transformers import LogitsWarper
 
 
-class AdvancedRepetitionPenaltyLogitsProcessor(LogitsProcessor):
+class AdvancedRepetitionPenaltyLogitsProcessor(LogitsWarper):
     def __init__(self, *args, **kwargs):
         pass