pull/889/merge
Youngju KIM 2026-10-17 03:58:09 +09:00 committed by GitHub
commit 91d2ef159b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 4 additions and 3 deletions

View File

@ -1,3 +1,4 @@
import os
import time import time
import argparse import argparse
import random import random
@ -11,7 +12,7 @@ warnings.filterwarnings('ignore')
def init_model(args): def init_model(args):
tokenizer = AutoTokenizer.from_pretrained(args.load_from) tokenizer = AutoTokenizer.from_pretrained(args.load_from)
if 'model' in args.load_from: if os.path.basename(os.path.normpath(args.load_from)) == 'model':
model = MiniMindForCausalLM(MiniMindConfig( model = MiniMindForCausalLM(MiniMindConfig(
hidden_size=args.hidden_size, hidden_size=args.hidden_size,
num_hidden_layers=args.num_hidden_layers, num_hidden_layers=args.num_hidden_layers,

View File

@ -56,7 +56,7 @@ TEST_CASES = [
def init_model(args): def init_model(args):
tokenizer = AutoTokenizer.from_pretrained(args.load_from) tokenizer = AutoTokenizer.from_pretrained(args.load_from)
if 'model' in args.load_from: if os.path.basename(os.path.normpath(args.load_from)) == 'model':
model = MiniMindForCausalLM(MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))) model = MiniMindForCausalLM(MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe)))
moe_suffix = '_moe' if args.use_moe else '' moe_suffix = '_moe' if args.use_moe else ''
ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'

View File

@ -27,7 +27,7 @@ app = FastAPI()
def init_model(args): def init_model(args):
tokenizer = AutoTokenizer.from_pretrained(args.load_from) tokenizer = AutoTokenizer.from_pretrained(args.load_from)
if 'model' in args.load_from: if os.path.basename(os.path.normpath(args.load_from)) == 'model':
moe_suffix = '_moe' if args.use_moe else '' moe_suffix = '_moe' if args.use_moe else ''
ckp = f'../{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' ckp = f'../{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
model = MiniMindForCausalLM(MiniMindConfig( model = MiniMindForCausalLM(MiniMindConfig(