import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# 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():
    model = AutoModelForCausalLM.from_pretrained(
        "anhvv200053/Vinallama-2-7B-updated1-instruction-v2",
        quantization_config=bnb_config,
        device_map={"": 0},
        use_auth_token=True
    )
    model.config.pretraining_tp = 1

    tokenizer = AutoTokenizer.from_pretrained('anhvv200053/Vinallama-2-7B-updated1-instruction-v2', trust_remote_code=True, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token

    return model, tokenizer