From 9e35e5348fe5ea223d2859ce2343fb4f0b2fb077 Mon Sep 17 00:00:00 2001
From: math-zhuxy <235934429@qq.com>
Date: Sun, 25 May 2025 16:14:33 +0800
Subject: [PATCH] new model_chat.py

---
 .gitignore    |   3 +-
 model_chat.py | 100 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 1 deletion(-)
 create mode 100644 model_chat.py

diff --git a/.gitignore b/.gitignore
index de6bd7a..ec6a5f5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
 /model/__pycache__
 /dataset
-/out
\ No newline at end of file
+/out
+/MiniMind2
\ No newline at end of file
diff --git a/model_chat.py b/model_chat.py
new file mode 100644
index 0000000..6a08f19
--- /dev/null
+++ b/model_chat.py
@@ -0,0 +1,100 @@
+import argparse
+import random
+import warnings
+import numpy as np
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.model_lora import *
+
+warnings.filterwarnings('ignore')
+
+
+def init_model(args):
+    transformers_model_path = './MiniMind2'
+    tokenizer = AutoTokenizer.from_pretrained(transformers_model_path)
+    model = AutoModelForCausalLM.from_pretrained(transformers_model_path, trust_remote_code=True)
+    print(f'MiniMind模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f}M(illion)')
+    return model.eval().to(args.device), tokenizer
+
+
+# 设置可复现的随机种子
+def setup_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def main():
+    print("开始加载模型")
+    parser = argparse.ArgumentParser(description="Chat with MiniMind")
+    parser.add_argument('--lora_name', default='None', type=str)
+    parser.add_argument('--out_dir', default='out', type=str)
+    parser.add_argument('--temperature', default=0.85, type=float)
+    parser.add_argument('--top_p', default=0.85, type=float)
+    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
+    # 此处max_seq_len（最大输出长度）并不意味模型具有对应的长文本的性能，仅防止QA出现被截断的问题
+    # MiniMind2-moe (145M)：(hidden_size=640, num_hidden_layers=8, use_moe=True)
+    # MiniMind2-Small (26M)：(hidden_size=512, num_hidden_layers=8)
+    # MiniMind2 (104M)：(hidden_size=768, num_hidden_layers=16)
+    parser.add_argument('--hidden_size', default=512, type=int)
+    parser.add_argument('--num_hidden_layers', default=8, type=int)
+    parser.add_argument('--max_seq_len', default=8192, type=int)
+    parser.add_argument('--use_moe', default=False, type=bool)
+    # 携带历史对话上下文条数
+    # history_cnt需要设为偶数，即【用户问题, 模型回答】为1组；设置为0时，即当前query不携带历史上文
+    # 模型未经过外推微调时，在更长的上下文的chat_template时难免出现性能的明显退化，因此需要注意此处设置
+    parser.add_argument('--history_cnt', default=0, type=int)
+    parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载")
+    parser.add_argument('--model_mode', default=1, type=int,
+                        help="0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型，4: RLAIF-Chat模型")
+    args = parser.parse_args()
+
+    model, tokenizer = init_model(args)
+
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    messages = []
+    for prompt in iter(lambda: input('👶: '), ''):
+        if prompt == "exit":
+            break
+        setup_seed(random.randint(0, 2048))
+
+        messages.append({"role": "user", "content": prompt})
+
+        new_prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        ) if args.model_mode != 0 else (tokenizer.bos_token + prompt)
+
+        inputs = tokenizer(
+            new_prompt,
+            return_tensors="pt",
+            truncation=True
+        ).to(args.device)
+
+        print('🤖️: ', end='')
+        generated_ids = model.generate(
+            inputs["input_ids"],
+            max_new_tokens=args.max_seq_len,
+            num_return_sequences=1,
+            do_sample=True,
+            attention_mask=inputs["attention_mask"],
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            streamer=streamer,
+            top_p=args.top_p,
+            temperature=args.temperature
+        )
+
+        response = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+        messages.append({"role": "assistant", "content": response})
+        print('\n')
+
+
+if __name__ == "__main__":
+    main()