Add 8-bit quantization for faster inference

app.py

@@ -4,18 +4,24 @@ HuggingFace Spaces 推理应用
 """
 
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
 
 # 模型配置
 MODEL_NAME = "zhman/llama-SFT-GRPO"
 
+# 配置 8-bit 量化以提升速度和减少内存占用
+quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_threshold=6.0
+)
+
 # 加载模型和分词器
 print("🔄 加载模型...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.bfloat16,
+    quantization_config=quantization_config,  # 使用 8-bit 量化
     device_map="auto"
 )
 print("✅ 模型加载完成！")