Update model.py

model.py

@@ -5,8 +5,28 @@ import os
 hf_token = os.getenv("LLM_token")
 os.environ["HUGGINGFACE_HUB_TOKEN"] = hf_token
 
-bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4")
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    llm_int8_skip_modules=None
+)
 
+def load_model(model_path="meta-llama/Meta-Llama-3-8B-Instruct"):
+    # 不要使用 torch.device 來傳入 pipeline，改用 device_map="auto"
+    pipe = pipeline(
+        "text-generation",
+        model=model_path,
+        model_kwargs={
+            "quantization_config": bnb_config,
+            "device_map": "auto",
+            "torch_dtype": torch.float16
+        },
+        token=hf_token
+    )
+    return pipe
+
+'''
 def load_model(model_path="meta-llama/Meta-Llama-3-8B-Instruct"):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -19,3 +39,5 @@ def load_model(model_path="meta-llama/Meta-Llama-3-8B-Instruct"):
         token=hf_token
     )
     return pipe
+'''
+