Add 4-bit BnB toggle

modeling/inference_models/generic_hf_torch/class.py

@@ -20,9 +20,29 @@ model_backend_name = "Huggingface"
 model_backend_type = "Huggingface" #This should be a generic name in case multiple model backends are compatible (think Hugging Face Custom and Basic Hugging Face)
 
 class model_backend(HFTorchInferenceModel):
+        
     def _initialize_model(self):
         return
 
+    def get_requested_parameters(self, model_name, model_path, menu_path, parameters = {}):
+        requested_parameters = super().get_requested_parameters(model_name, model_path, menu_path, parameters = {})
+        requested_parameters.append({
+                                    "uitype": "toggle",
+                                    "unit": "bool",
+                                    "label": "Use 4-bit",
+                                    "id": "use_4_bit",
+                                    "default": False,
+                                    "tooltip": "Whether or not to use BnB's 4-bit mode",
+                                    "menu_path": "Layers",
+                                    "extra_classes": "",
+                                    "refresh_model_inputs": False
+                                })
+        return requested_parameters
+ 
+    def set_input_parameters(self, parameters):
+        super().set_input_parameters(parameters)
+        self.use_4_bit = parameters['use_4_bit']
+
     def _load(self, save_model: bool, initial_load: bool) -> None:
         utils.koboldai_vars.allowsp = True
 
@@ -51,6 +71,12 @@ class model_backend(HFTorchInferenceModel):
             "low_cpu_mem_usage": True,
         }
         
+        if self.use_4_bit:
+            self.lazy_load = False
+            tf_kwargs.update({
+                "load_in_4bit": True,
+            })
+
         if self.model_type == "gpt2":
             # We must disable low_cpu_mem_usage and if using a GPT-2 model
             # because GPT-2 is not compatible with this feature yet.