[feat] update import

pull/825/head
jingyaogong 2025-10-31 23:45:55 +08:00
parent fdc6c5db5d
commit 7c3d660eec
1 changed files with 10 additions and 11 deletions

View File

@ -8,6 +8,8 @@ import numpy as np
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from torch.utils.data import Sampler from torch.utils.data import Sampler
from transformers import AutoTokenizer
from model.model_minimind import MiniMindForCausalLM
def is_main_process(): def is_main_process():
@ -96,9 +98,6 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device='cuda'): def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device='cuda'):
from transformers import AutoTokenizer
from model.model_minimind import MiniMindForCausalLM
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
model = MiniMindForCausalLM(lm_config) model = MiniMindForCausalLM(lm_config)