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(
#     load_in_4bit=True,
#     bnb_4bit_quant_type='nf4',
#     bnb_4bit_compute_dtype="float16",
#     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/Vinallam_vssai_model",
        device_map={"": "cpu"},  # Đặt sử dụng CPU
        token = token
    )
    # model.config.pretraining_tp = 1

    tokenizer = AutoTokenizer.from_pretrained('anhvv200053/Vinallam_vssai_model', trust_remote_code=True, use_fast=True, token = token)
    tokenizer.pad_token = tokenizer.eos_token

    return model, tokenizer