[feat] update import
parent
fdc6c5db5d
commit
7c3d660eec
|
|
@ -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)
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue