Update app.py

app.py

@@ -2,7 +2,7 @@ import subprocess
 import sys
 
 # Cài đặt các thư viện nếu chưa có
-subprocess.check_call([sys.executable, "-m", "pip", "install", "transformers", "streamlit", "torch", "peft"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "transformers", "streamlit", "torch", "bitsandbytes","peft"])
 
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
@@ -13,20 +13,15 @@ import torch
 BASE_MODEL_NAME = "unsloth/deepseek-r1-distill-llama-8b-unsloth-bnb-4bit"
 ADAPTER_MODEL_PATH = "lora_model"
 
-# Load mô hình gốc trên CPU
-base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL_NAME,
-    torch_dtype=torch.float32,  # sử dụng float32 cho CPU
-    device_map="cpu"            # ép chạy trên CPU
-)
+# Load mô hình gốc
+base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
 # Áp dụng adapter LoRA
 model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL_PATH)
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
 
 def generate_response(prompt):
     """Generate a response from the model."""
-    # Chuyển dữ liệu input sang CPU
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
     streamer = TextStreamer(tokenizer)
     with torch.no_grad():
         model.generate(**inputs, streamer=streamer, max_length=512)
@@ -59,4 +54,4 @@ if user_input:
         st.markdown(response)
 
     # Append assistant response
-    st.session_state.messages.append({"role": "assistant", "content": response})
+    st.session_state.messages.append({"role": "assistant", "content": response})