ncc

app.py

@@ -2,14 +2,53 @@
 
 transformers 4.31.0
 """
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import os
+import time
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
 from transformers.generation import GenerationConfig
+from transformers import BitsAndBytesConfig
+from loguru import logger
 
-tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-7B-Chat", trust_remote_code=True)
+os.environ["TZ"] = "Asia/Shanghai"
+try:
+    time.tzset()  # type: ignore # pylint: disable=no-member
+except Exception:
+    # Windows
+    logger.warning("Windows, cant run time.tzset()")
 
-model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B-Chat", device_map="auto", trust_remote_code=True).eval()
-model.generation_config = GenerationConfig.from_pretrained("Qwen/Qwen-7B-Chat", trust_remote_code=True) # 可指定不同的生成长度、top_p等相关超参
+device_map = "cuda:0" if torch.cuda.is_available() else "cpu"
+# has_cuda = False  # force cpu
 
-response, history = model.chat(tokenizer, "你好", history=None)
-# response, history = model.chat(tokenizer, "你好", history=[])
+model_name = "Qwen/Qwen-7B-Chat"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+# quantization configuration for NF4 (4 bits)
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type='nf4',
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+# quantization configuration for Int8 (8 bits)
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map=device_map,
+    quantization_config=quantization_config,
+    # max_memory=max_memory,
+    trust_remote_code=True,
+).eval()
+
+# model = model.eval()
+
+# model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B-Chat", device_map="auto", trust_remote_code=True).eval()
+
+# 可指定不同的生成长度、top_p等相关超参
+model.generation_config = GenerationConfig.from_pretrained("Qwen/Qwen-7B-Chat", trust_remote_code=True) 
+
+# response, history = model.chat(tokenizer, "你好", history=None)
+response, history = model.chat(tokenizer, "你好", history=[])
 print(response)
+

requirements.txt

@@ -11,3 +11,7 @@ einops
 
 torch  # 2.0.1
 safetensors
+bitsandbytes
+
+loguru
+about-time