Update utils.py

utils.py

@@ -22,6 +22,7 @@ import transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM, GPT2Tokenizer, GPT2LMHeadModel
 #import auto_gptq
 #from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
+from transformers import LlamaForCausalLM, LlamaTokenizer
 
 
 def reset_state():
@@ -99,6 +100,19 @@ def load_tokenizer_and_model(base_model,load_8bit=False):
     return tokenizer,model,device
 
 
+def load_tokenizer_and_model_Baize(base_model, load_8bit=True):
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+    
+    
+    tokenizer = LlamaTokenizer.from_pretrained(base_model, add_eos_token=True, use_auth_token=True)
+    model = LlamaForCausalLM.from_pretrained(base_model, load_in_8bit=True, device_map="auto")
+
+    return tokenizer,model, device
+
+
 def load_tokenizer_and_model_gpt2(base_model,load_8bit=False):
     if torch.cuda.is_available():
         device = "cuda"