obscura/minimind
1
0
Fork 0
mirror of https://github.com/jingyaogong/minimind.git synced 2026-10-18 23:07:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

[feat] add args

Browse Source
This commit is contained in:
jingyaogong
2025-10-30 10:05:12 +08:00
parent 800fed4639
commit bf123b585d
11 changed files with 57 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/serve_openai_api.py
+22 -23
View File
@@ -25,26 +25,24 @@ app = FastAPI()
def init_model(args):
if args.load == 0:
tokenizer = AutoTokenizer.from_pretrained('../model/')
moe_path = '_moe' if args.use_moe else ''
modes = {0: 'pretrain', 1: 'full_sft', 2: 'rlhf', 3: 'reason'}
ckp = f'../{args.out_dir}/{modes[args.model_mode]}_{args.hidden_size}{moe_path}.pth'
tokenizer = AutoTokenizer.from_pretrained('../model/')
if args.load_from == 'model':
moe_suffix = '_moe' if args.use_moe else ''
ckp = f'../{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
model = MiniMindForCausalLM(MiniMindConfig(
hidden_size=args.hidden_size,
num_hidden_layers=args.num_hidden_layers,
max_seq_len=args.max_seq_len,
use_moe=args.use_moe
use_moe=bool(args.use_moe)
))
model.load_state_dict(torch.load(ckp, map_location=device), strict=True)
if args.lora_name != 'None':
if args.lora_weight != 'None':
apply_lora(model)
load_lora(model, f'../{args.out_dir}/{args.lora_name}_{args.hidden_size}.pth')
load_lora(model, f'../{args.save_dir}/lora/{args.lora_weight}_{args.hidden_size}.pth')
else:
model_path = '../MiniMind2'
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_path)
print(f'MiniMind模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f}M(illion)')
model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(args.load_from)
print(f'MiniMind模型参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M(illion)')
return model.eval().to(device), tokenizer
@@ -163,15 +161,16 @@ async def chat_completions(request: ChatRequest):
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Server for MiniMind")
parser.add_argument('--out_dir', default='out', type=str)
parser.add_argument('--lora_name', default='None', type=str)
parser.add_argument('--hidden_size', default=768, type=int)
parser.add_argument('--num_hidden_layers', default=16, type=int)
parser.add_argument('--max_seq_len', default=8192, type=int)
parser.add_argument('--use_moe', default=False, type=bool)
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模型")
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model, tokenizer = init_model(parser.parse_args())
parser.add_argument('--load_from', default='model', type=str, help="模型加载路径（model=原生torch权重，其他路径=transformers格式）")
parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀（pretrain, full_sft, dpo, reason, ppo_actor, grpo, spo）")
parser.add_argument('--lora_weight', default='None', type=str, help="LoRA权重名称（None表示不使用，可选：lora_identity, lora_medical）")
parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度（512=Small-26M, 640=MoE-145M, 768=Base-104M）")
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量（Small/MoE=8, Base=16）")
parser.add_argument('--max_seq_len', default=8192, type=int, help="最大序列长度")
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
args = parser.parse_args()
device = args.device
model, tokenizer = init_model(args)
uvicorn.run(app, host="0.0.0.0", port=8998)