test app

app.py

@@ -36,40 +36,21 @@ from transformers.generation import GenerationConfig
 DEFAULT_CKPT_PATH = 'WisdomShell/CodeShell-7B-Chat'
 
 
-def _get_args():
-    parser = ArgumentParser()
-    parser.add_argument("-c", "--checkpoint-path", type=str, default=DEFAULT_CKPT_PATH,
-                        help="Checkpoint name or path, default to %(default)r")
-    parser.add_argument("--device", type=str, default="cuda:0", help="GPU device.")
-
-    parser.add_argument("--share", action="store_true", default=False,
-                        help="Create a publicly shareable link for the interface.")
-    parser.add_argument("--inbrowser", action="store_true", default=False,
-                        help="Automatically launch the interface in a new tab on the default browser.")
-    parser.add_argument("--server-port", type=int, default=8000,
-                        help="Demo server port.")
-    parser.add_argument("--server-name", type=str, default="127.0.0.1",
-                        help="Demo server name.")
-
-    args = parser.parse_args()
-    return args
-
-
 def _load_model_tokenizer(args):
     tokenizer = AutoTokenizer.from_pretrained(
-        args.checkpoint_path, trust_remote_code=True, resume_download=True,
+        'WisdomShell/CodeShell-7B-Chat', trust_remote_code=True, resume_download=True,
     )
 
     model = AutoModelForCausalLM.from_pretrained(
-        args.checkpoint_path,
-        device_map=args.device,
+        'WisdomShell/CodeShell-7B-Chat',
+        device_map='cpu',
         trust_remote_code=True,
         resume_download=True,
         torch_dtype=torch.bfloat16
     ).eval()
 
     config = GenerationConfig.from_pretrained(
-        args.checkpoint_path, trust_remote_code=True, resume_download=True,
+        'WisdomShell/CodeShell-7B-Chat', trust_remote_code=True, resume_download=True,
     )
 
     return model, tokenizer, config
@@ -188,21 +169,12 @@ including hate speech, violence, pornography, deception, etc. \
 (注：本演示受CodeShell的许可协议限制。我们强烈建议，用户不应传播及不应允许他人传播以下内容，\
 包括但不限于仇恨言论、暴力、色情、欺诈相关的有害信息。)""")
 
-    demo.queue().launch(
-        share=args.share,
-        inbrowser=args.inbrowser,
-        server_port=args.server_port,
-        server_name=args.server_name,
-    )
-
+    demo.queue().launch()
 
-def main():
-    args = _get_args()
 
-    model, tokenizer, config = _load_model_tokenizer(args)
+args = {}
 
-    _launch_demo(args, model, tokenizer, config)
+model, tokenizer, config = _load_model_tokenizer(args)
 
+_launch_demo(args, model, tokenizer, config)
 
-if __name__ == '__main__':
-    main()