Update utils.py

utils.py

@@ -95,6 +95,34 @@ def load_tokenizer_and_model(base_model,load_8bit=False):
     model.eval()
     return tokenizer,model,device
 
+
+def load_tokenizer_and_model_gpt2(base_model,load_8bit=False):
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+
+    tokenizer = GPT2Tokenizer.from_pretrained(base_model, use_fast = True)
+    if device == "cuda":
+        model = GPT2LMHeadModel.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 tokenizer,model,device
+
+
+
 # Greedy Search
 def greedy_search(input_ids: torch.Tensor,
                   model: torch.nn.Module,