Better working valid detection and named model backends for UI

modeling/inference_models/api.py

@@ -17,12 +17,13 @@ from modeling.inference_model import (
     ModelCapabilities,
 )
 
+model_backend_name = "KoboldAI API"
 
 class APIException(Exception):
     """To be used for errors when using the Kobold API as an interface."""
 
 
-class model_loader(InferenceModel):
+class model_backend(InferenceModel):
     def __init__(self) -> None:
         super().__init__()
         #self.base_url = ""

modeling/inference_models/basic_api.py

@@ -15,11 +15,13 @@ from modeling.inference_model import (
 )
 
 
+model_backend_name = "KoboldAI Old Colab Method"
+
 class BasicAPIException(Exception):
     """To be used for errors when using the Basic API as an interface."""
 
 
-class model_loader(InferenceModel):
+class model_backend(InferenceModel):
     def __init__(self) -> None:
         super().__init__()
 

modeling/inference_models/generic_hf_torch.py

@@ -22,8 +22,9 @@ except ModuleNotFoundError as e:
 
 from modeling.inference_models.parents.hf_torch import HFTorchInferenceModel
 
+model_backend_name = "Huggingface"
 
-class model_loader(HFTorchInferenceModel):
+class model_backend(HFTorchInferenceModel):
     
     def _initialize_model(self):
         return

modeling/inference_models/gooseai.py

@@ -11,16 +11,17 @@ from modeling.inference_model import (
     InferenceModel,
 )
 
-from modeling.inference_models.parents.openai_gooseai import model_loader as openai_gooseai_model_loader
+from modeling.inference_models.parents.openai_gooseai import model_backend as openai_gooseai_model_backend
 
 
+model_backend_name = "GooseAI"
 
 class OpenAIAPIError(Exception):
     def __init__(self, error_type: str, error_message) -> None:
         super().__init__(f"{error_type}: {error_message}")
 
 
-class model_loader(openai_gooseai_model_loader):
+class model_backend(openai_gooseai_model_backend):
     """InferenceModel for interfacing with OpenAI's generation API."""
     
     def __init__(self):

modeling/inference_models/hf_mtj.py

@@ -19,10 +19,10 @@ from modeling.inference_model import (
 from modeling.inference_models.parents.hf import HFInferenceModel
 from modeling.tokenizer import GenericTokenizer
 
+model_backend_name = "Huggingface MTJ"
 
 
-
-class model_loader(HFInferenceModel):
+class model_backend(HFInferenceModel):
     def __init__(
         self,
         #model_name: str,

modeling/inference_models/horde.py

@@ -16,12 +16,13 @@ from modeling.inference_model import (
     ModelCapabilities,
 )
 
+model_backend_name = "Horde"
 
 class HordeException(Exception):
     """To be used for errors on server side of the Horde."""
 
 
-class model_loader(InferenceModel):
+class model_backend(InferenceModel):
     def __init__(self) -> None:
         super().__init__()
         self.url = "https://horde.koboldai.net"

modeling/inference_models/openai.py

@@ -11,16 +11,16 @@ from modeling.inference_model import (
     InferenceModel,
 )
 
-from modeling.inference_models.parents.openai_gooseai import model_loader as openai_gooseai_model_loader
-
+from modeling.inference_models.parents.openai_gooseai import model_backend as openai_gooseai_model_backend
 
+model_backend_name = "OpenAI"
 
 class OpenAIAPIError(Exception):
     def __init__(self, error_type: str, error_message) -> None:
         super().__init__(f"{error_type}: {error_message}")
 
 
-class model_loader(openai_gooseai_model_loader):
+class model_backend(openai_gooseai_model_backend):
     """InferenceModel for interfacing with OpenAI's generation API."""
     
     def __init__(self):

modeling/inference_models/parents/hf.py

@@ -1,7 +1,7 @@
 import os
 from typing import Optional
 from transformers import AutoConfig
-
+import warnings
 import utils
 import koboldai_settings
 from logger import logger
@@ -43,7 +43,7 @@ class HFInferenceModel(InferenceModel):
         else:
             self.model_config = AutoConfig.from_pretrained(model_name, revision=utils.koboldai_vars.revision, cache_dir="cache")
         layer_count = self.model_config["n_layer"] if isinstance(self.model_config, dict) else self.model_config.num_layers if hasattr(self.model_config, "num_layers") else self.model_config.n_layer if hasattr(self.model_config, "n_layer") else self.model_config.num_hidden_layers if hasattr(self.model_config, 'num_hidden_layers') else None
-        if layer_count is not None and layer_count >= 0:
+        if layer_count is not None and layer_count >= 0 and not self.nobreakmodel:
             if os.path.exists("settings/{}.breakmodel".format(model_name.replace("/", "_"))):
                 with open("settings/{}.breakmodel".format(model_name.replace("/", "_")), "r") as file:
                     data = [x for x in file.read().split("\n")[:2] if x != '']
@@ -128,15 +128,17 @@ class HFInferenceModel(InferenceModel):
     def set_input_parameters(self, parameters):
         if self.hf_torch:
             import breakmodel
-            gpu_count = torch.cuda.device_count()
-            layers = []
-            for i in range(gpu_count):
-                layers.append(int(parameters["{}_Layers".format(i)]) if parameters["{}_Layers".format(i)].isnumeric() else None)
-            self.cpu_layers = parameters['CPU_Layers'] if 'CPU_Layers' in parameters else None
-            self.layers = layers
-            self.disk_layers = int(parameters['Disk_Layers']) if 'Disk_Layers' in parameters and parameters['Disk_Layers'].isnumeric() else 0    
-            breakmodel.gpu_blocks = layers
-            breakmodel.disk_blocks = self.disk_layers
+            layer_count = self.model_config["n_layer"] if isinstance(self.model_config, dict) else self.model_config.num_layers if hasattr(self.model_config, "num_layers") else self.model_config.n_layer if hasattr(self.model_config, "n_layer") else self.model_config.num_hidden_layers if hasattr(self.model_config, 'num_hidden_layers') else None
+            if layer_count is not None and layer_count >= 0 and not self.nobreakmodel:
+                gpu_count = torch.cuda.device_count()
+                layers = []
+                for i in range(gpu_count):
+                    layers.append(int(parameters["{}_Layers".format(i)]) if parameters["{}_Layers".format(i)].isnumeric() else None)
+                self.cpu_layers = parameters['CPU_Layers'] if 'CPU_Layers' in parameters else None
+                self.layers = layers
+                self.disk_layers = int(parameters['Disk_Layers']) if 'Disk_Layers' in parameters and parameters['Disk_Layers'].isnumeric() else 0    
+                breakmodel.gpu_blocks = layers
+                breakmodel.disk_blocks = self.disk_layers
             self.usegpu = parameters['use_gpu'] if 'use_gpu' in parameters else None
             self.model_type = self.get_model_type()
             self.breakmodel = ((self.model_type != 'gpt2') or self.model_type in ("gpt_neo", "gptj", "xglm", "opt")) and not self.nobreakmodel

modeling/inference_models/parents/openai_gooseai.py

@@ -18,7 +18,7 @@ class OpenAIAPIError(Exception):
         super().__init__(f"{error_type}: {error_message}")
 
 
-class model_loader(InferenceModel):
+class model_backend(InferenceModel):
     """InferenceModel for interfacing with OpenAI's generation API."""
     
     def __init__(self):

modeling/inference_models/readonly.py

@@ -14,12 +14,13 @@ from modeling.inference_model import (
     ModelCapabilities,
 )
 
+model_backend_name = "Read Only"
 
 class BasicAPIException(Exception):
     """To be used for errors when using the Basic API as an interface."""
 
 
-class model_loader(InferenceModel):
+class model_backend(InferenceModel):
     def __init__(self) -> None:
         super().__init__()
 

modeling/inference_models/rwkv.py

@@ -55,7 +55,10 @@ MODEL_FILES = {
 }
 
 
-class model_loader(InferenceModel):
+model_backend_name = "RWKV"
+
+
+class model_backend(InferenceModel):
     def __init__(
         self,
         #model_name: str,