Update app.py

app.py

@@ -1,18 +1,45 @@
 import os 
+import time
+
 # os.system("pip install --upgrade torch transformers sentencepiece scipy cpm_kernels accelerate bitsandbytes loguru")
-os.system("pip install transformers loguru")
+os.system("pip install torch transformers sentencepiece loguru")
 
 import gradio as gr
 from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
 
+# fix timezone in Linux
+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_name = "THUDM/chatglm2-6b-int4"  # 3.9G
+
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
+
+has_cuda = torch.cuda.is_available()
+# has_cuda = False  # force cpu
+
 logger.debug("load")
-model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
+if has_cuda:
+    if model_name.endswith("int4"):
+        model = AutoModel.from_pretrained(model_name, trust_remote_code=True).cuda()
+    else:
+        model = (
+            AutoModel.from_pretrained(model_name, trust_remote_code=True).cuda().half()
+        )
+else:
+    model = AutoModel.from_pretrained(
+        model_name, trust_remote_code=True
+    ).half()  # .float() .half().float()
+
+model = model.eval()
 logger.debug("done load")
+
 # tokenizer = AutoTokenizer.from_pretrained("openchat/openchat_v2_w")
 # model = AutoModelForCausalLM.from_pretrained("openchat/openchat_v2_w", load_in_8bit_fp32_cpu_offload=True, load_in_8bit=True)
-model.half()
-model = model.eval()
 
 model_path = model.config._dict['model_name_or_path']
 logger.debug(f"{model_path=}")