Update model_setup.py

model_setup.py

@@ -1,6 +1,8 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import os
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM
 
 # Cấu hình BitsAndBytes để tải mô hình 4-bit
 bnb_config = BitsAndBytesConfig(
@@ -10,15 +12,22 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=False,
 )
 
+
+
 # Thiết lập mô hình và tokenizer
 def load_model():
     token = os.getenv("HF_TOKEN")  # Lấy token từ biến môi trường
+
+
+    config = PeftConfig.from_pretrained("anhvv200053/Vinallama-2-7B-updated1-instruction-v2")
+    base_model = AutoModelForCausalLM.from_pretrained("vilm/vinallama-2.7b-chat")
+    model = PeftModel.from_pretrained(base_model, "anhvv200053/Vinallama-2-7B-updated1-instruction-v2", token = token)
     
-    model = AutoModelForCausalLM.from_pretrained(
-        "anhvv200053/Vinallama-2-7B-updated1-instruction-v2",
-        device_map={"": 0},
-        token = token
-    )
+    # model = AutoModelForCausalLM.from_pretrained(
+    #     "anhvv200053/Vinallama-2-7B-updated1-instruction-v2",
+    #     device_map={"": 0},
+    #     token = token
+    # )
     model.config.pretraining_tp = 1
 
     tokenizer = AutoTokenizer.from_pretrained('anhvv200053/Vinallama-2-7B-updated1-instruction-v2', trust_remote_code=True, use_fast=True, token = token)