Update app.py

app.py

@@ -1,23 +1,31 @@
 import os
 from huggingface_hub import login
-login(token=os.environ["HF_TOKEN"])  # Dùng biến môi trường để lấy token 
+login(token=os.environ["HF_TOKEN"])  # Dùng biến môi trường để lấy token
 
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import gradio as gr
 
 # Load model
 model_id = "MindVR/JohnTran_Fine-tune"
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=True)
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ["HF_TOKEN"])
+
+# Thiết lập 4bit config
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+)
+
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype=torch.float16,
-    device_map="auto",             # Tự động GPU
-    low_cpu_mem_usage=True,        # Tối ưu RAM
-    token=os.environ["HF_TOKEN"]   # ✅ thay vì use_auth_token=True
+    device_map="auto",
+    quantization_config=bnb_config,
+    low_cpu_mem_usage=True,
+    token=os.environ["HF_TOKEN"]
 )
 
-
 # Hàm xử lý yêu cầu
 def chat(prompt):
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")