Update app.py

app.py

@@ -3,10 +3,12 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
+import bitsandbytes as bnb
+
 # 設置模型 ID 和加載 Hugging Face API token
 model_id = "dingckc/FineLlama-3.1-8B"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.getenv('ACCESS_KEY'))
-model = AutoModelForCausalLM.from_pretrained(model_id, token=os.getenv('ACCESS_KEY'))
+model = AutoModelForCausalLM.from_pretrained(model_id, load_in_8bit=True, token=os.getenv('ACCESS_KEY'))
 model = model.to("cuda" if torch.cuda.is_available() else "cpu")
 
 # 定義推理函數