Update app.py

app.py

@@ -1,14 +1,14 @@
-from loguru import logger
-import rich
+# pylint: disable=invalid-name, line-too-long, missing-module-docstring
+import gc
 import os
 import time
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import gc
 
-from huggingface_hub import hf_hub_download
+import rich
+import torch
 from huggingface_hub import snapshot_download
-# snapshot_download(repo_id="lysandre/arxiv-nlp")
+from loguru import logger
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.generation.utils import GenerationConfig
 
 model_name = "baichuan-inc/Baichuan2-13B-Chat-4bits"
 # snapshot_download?
@@ -18,7 +18,7 @@ loc = snapshot_download(repo_id=model_name, local_dir="model")
 os.environ["TZ"] = "Asia/Shanghai"
 try:
     time.tzset()  # type: ignore # pylint: disable=no-member
-except Exception:
+except Exception:  # pylint: disable=broad-except
     # Windows
     logger.warning("Windows, cant run time.tzset()")
 
@@ -30,16 +30,19 @@ has_cuda = torch.cuda.is_available()
 
 if has_cuda:
     model = AutoModelForCausalLM.from_pretrained(
-        "model",   # loc
-        # device_map="auto", 
-        torch_dtype=torch.bfloat16, 
+        "model",  # loc
+        # device_map="auto",
+        torch_dtype=torch.bfloat16,  # pylint: disable=no-member
         load_in_8bit=True,
         trust_remote_code=True,
         # use_ram_optimized_load=False,
         # offload_folder="offload_folder",
     ).cuda()
 else:
-    model = AutoModel.from_pretrained(model_name, trust_remote_code=True).float()
+    # model = AutoModel.from_pretrained(model_name, trust_remote_code=True).float()
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, trust_remote_code=True
+    ).float()
 
 model = model.eval()
 
@@ -47,18 +50,19 @@ rich.print(f"{model=}")
 
 logger.info("done")
 
-# ========
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers.generation.utils import GenerationConfig
-tokenizer = AutoTokenizer.from_pretrained("baichuan-inc/Baichuan2-13B-Chat-4bits", use_fast=False, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(
+    "baichuan-inc/Baichuan2-13B-Chat-4bits", use_fast=False, trust_remote_code=True
+)
 
 # model = AutoModelForCausalLM.from_pretrained("baichuan-inc/Baichuan2-13B-Chat-4bits", device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
 
-model.generation_config = GenerationConfig.from_pretrained("baichuan-inc/Baichuan2-13B-Chat-4bits")
+model.generation_config = GenerationConfig.from_pretrained(
+    "baichuan-inc/Baichuan2-13B-Chat-4bits"
+)
 messages = []
 messages.append({"role": "user", "content": "解释一下“温故而知新”"})
 response = model.chat(tokenizer, messages)
 
 rich.print(response)
 
-logger.info(f"{response=}")
+logger.info(f"{response=}")