Specify torch_type and device_map when loading the model

Author: underspirit

text_generation_demo.py

@@ -10,8 +10,8 @@
 
 def init_model(args):
     global tokenizer, model
-    tokenizer = AutoTokenizer.from_pretrained(args.model_path, truncation_side="left", padding_side="left")
-    model = AutoModelForCausalLM.from_pretrained(args.tokenizer_path, trust_remote_code=True).half().cuda()
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, truncation_side="left", padding_side="left")
+    model = AutoModelForCausalLM.from_pretrained(args.model_path, trust_remote_code=True, torch_dtype=torch.float16, device_map='auto')
     model = model.eval()
 
 
@@ -79,4 +79,4 @@ def get_args():
                                  temperature, top_k, top_p],
                          outputs=outputs)
 
-    demo.launch(server_name="0.0.0.0", server_port=args.port)
+    demo.launch(server_name="0.0.0.0", server_port=args.port)