Update huggingface_api.py (#2409)

fastchat/serve/huggingface_api.py

@@ -14,6 +14,7 @@
 
 @torch.inference_mode()
 def main(args):
+    # Load model
     model, tokenizer = load_model(
         args.model_path,
         device=args.device,
@@ -25,13 +26,14 @@ def main(args):
         debug=args.debug,
     )
 
+    # Build the prompt with a conversation template
     msg = args.message
-
     conv = get_conversation_template(args.model_path)
     conv.append_message(conv.roles[0], msg)
     conv.append_message(conv.roles[1], None)
     prompt = conv.get_prompt()
 
+    # Run inference
     inputs = tokenizer([prompt], return_tensors="pt").to(args.device)
     output_ids = model.generate(
         **inputs,
@@ -49,6 +51,7 @@ def main(args):
         output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
     )
 
+    # Print results
     print(f"{conv.roles[0]}: {msg}")
     print(f"{conv.roles[1]}: {outputs}")