Update main.py

main.py

@@ -13,43 +13,43 @@ People will talk to you about their personal life/mental health issues and you w
 therapist would, while being empathetic."""
 
 
-# def load_peft_model_and_tokenizer(peft_model, base_model):
+def load_peft_model_and_tokenizer(peft_model, base_model):
 
-#     tokenizer = AutoTokenizer.from_pretrained(peft_model)
+    tokenizer = AutoTokenizer.from_pretrained(peft_model)
 
 
-#     bnb_config = BitsAndBytesConfig(
-#         load_in_4bit=True,
-#         bnb_4bit_use_double_quant=True,
-#         bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
-#     )
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
+    )
 
-#    # Load base model
-#     base_model = AutoModelForCausalLM.from_pretrained(
-#         base_model,
-#         device_map="auto",
-#         torch_dtype=torch.bfloat16,
-#         quantization_config=bnb_config
-#     )
+   # Load base model
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        quantization_config=bnb_config
+    )
 
 
-#     base_model.resize_token_embeddings(len(tokenizer))
-#     model = PeftModel.from_pretrained(model=base_model, model_id=peft_model)
+    base_model.resize_token_embeddings(len(tokenizer))
+    model = PeftModel.from_pretrained(model=base_model, model_id=peft_model)
 
-#     return tokenizer, model
+    return tokenizer, model
 
-def load_peft_model_and_tokenizer(peft_model):
-    # Load the tokenizer from the specified model path or identifier
-    tokenizer = AutoTokenizer.from_pretrained(peft_model)
+# def load_peft_model_and_tokenizer(peft_model):
+#     # Load the tokenizer from the specified model path or identifier
+#     tokenizer = AutoTokenizer.from_pretrained(peft_model)
     
-    # Load the PEFT model for causal language modeling with specific device map and torch dtype
-    model = AutoPeftModelForCausalLM.from_pretrained(
-        peft_model,
-        device_map="auto",
-        torch_dtype=torch.float16
-    )
+#     # Load the PEFT model for causal language modeling with specific device map and torch dtype
+#     model = AutoPeftModelForCausalLM.from_pretrained(
+#         peft_model,
+#         device_map="auto",
+#         torch_dtype=torch.float16
+#     )
     
-    return tokenizer, model
+#     return tokenizer, model