Update utils.py

utils.py

@@ -79,6 +79,33 @@ def load_tokenizer_and_model(base_model, load_8bit=False):
     return tokenizer,model, device
 
 
+def load_model(base_model, load_8bit=False):
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+
+    if device == "cuda":
+        model = AutoModelForCausalLM.from_pretrained(
+            base_model,
+            load_in_8bit=load_8bit,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            base_model, device_map={"": device}, low_cpu_mem_usage=True
+        )
+
+    #if not load_8bit:
+        #model.half()  # seems to fix bugs for some users.
+
+    model.eval()
+    return model, device
+
+    
+
+
 def load_tokenizer(base_model):
     if torch.cuda.is_available():
         device = "cuda"