Initial commit

src/cli_demo.py

@@ -0,0 +1,66 @@
+# coding=utf-8
+# Implements stream chat in command line for LLaMA fine-tuned with PEFT.
+# Usage: python cli_demo.py --checkpoint_dir path_to_checkpoint
+
+
+import torch
+from utils import ModelArguments, auto_configure_device_map, load_pretrained
+from transformers import HfArgumentParser
+
+
+def main():
+
+    parser = HfArgumentParser(ModelArguments)
+    model_args, = parser.parse_args_into_dataclasses()
+    model, tokenizer = load_pretrained(model_args)
+    if torch.cuda.device_count() > 1:
+        from accelerate import dispatch_model
+        device_map = auto_configure_device_map(torch.cuda.device_count())
+        model = dispatch_model(model, device_map)
+    else:
+        model = model.cuda()
+    model.eval()
+
+    def predict(query, history: list):
+        inputs = tokenizer([query], return_tensors="pt")
+        inputs = inputs.to(model.device)
+        gen_kwargs = {
+            "do_sample": True,
+            "top_p": 0.9,
+            "top_k": 40,
+            "temperature": 0.7,
+            "num_beams": 1,
+            "max_new_tokens": 256,
+            "repetition_penalty": 1.5
+        }
+        with torch.no_grad():
+            generation_output = model.generate(**inputs, **gen_kwargs)
+        outputs = generation_output.tolist()[0][len(inputs["input_ids"][0]):]
+        response = tokenizer.decode(outputs, skip_special_tokens=True)
+        history = history + [(query, response)]
+        return response, history
+
+    history = []
+    print("欢迎使用 LLaMA-7B 模型，输入内容即可对话，clear清空对话历史，stop终止程序")
+    while True:
+        try:
+            query = input("\nInput: ")
+        except UnicodeDecodeError:
+            print("Detected decoding error at the inputs, please set the terminal encoding to utf-8.")
+            continue
+        except Exception:
+            raise
+
+        if query.strip() == "stop":
+            break
+
+        if query.strip() == "clear":
+            history = []
+            continue
+
+        response, history = predict(query, history)
+        print("LLaMA-7B:", response)
+
+
+if __name__ == "__main__":
+    main()