wehub-resource-sync
105 lines
4.2 KiB
Python
105 lines
4.2 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import json
|
|
import os
|
|
import shutil
|
|
import time
|
|
from typing import List, Optional, Union
|
|
|
|
from swift.arguments import SamplingArguments
|
|
from swift.dataset import load_dataset
|
|
from swift.utils import get_logger
|
|
from ..base import SwiftPipeline
|
|
from .distill_sampler import DistillSampler
|
|
from .vanilla_sampler import VanillaSampler
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class SwiftSampling(SwiftPipeline):
|
|
args_class = SamplingArguments
|
|
args: args_class
|
|
|
|
def __init__(self, args: Optional[Union[List[str], SamplingArguments]] = None) -> None:
|
|
super().__init__(args)
|
|
self.args.save_args()
|
|
os.makedirs(self.args.output_dir, exist_ok=True)
|
|
self.cur_piece = 0
|
|
self.total_piece = 1
|
|
|
|
if self.args.data_range:
|
|
self.cur_piece, self.total_piece = self.args.data_range
|
|
|
|
if self.args.sampler_type == 'sample':
|
|
self.sampler = VanillaSampler(self.args)
|
|
elif self.args.sampler_type == 'distill':
|
|
self.sampler = DistillSampler(self.args)
|
|
else:
|
|
raise ValueError(f'Unsupported sampler type: {self.args.sampler_type}')
|
|
|
|
def _get_dataset(self):
|
|
args = self.args
|
|
dataset_kwargs = args.get_dataset_kwargs()
|
|
sampling_dataset, _ = load_dataset(
|
|
args.dataset, split_dataset_ratio=0., shuffle=args.dataset_shuffle, **dataset_kwargs)
|
|
logger.info(f'Sampling_dataset: {sampling_dataset}')
|
|
dataset_len = len(sampling_dataset)
|
|
piece_len = dataset_len // self.total_piece
|
|
sampling_dataset = sampling_dataset.select(range(piece_len * self.cur_piece, piece_len * (self.cur_piece + 1)))
|
|
return sampling_dataset
|
|
|
|
def run(self):
|
|
os.makedirs(self.args.output_dir, exist_ok=True)
|
|
iter_file = os.path.join(self.args.output_dir, self.args.output_file)
|
|
resume_file = os.path.join(self.args.output_dir, self.args.output_file + '.resume')
|
|
tmp_file = os.path.join(self.args.output_dir, self.args.output_file + '.tmp')
|
|
ckpt_state_file = os.path.join(self.args.output_dir, 'ckpt_state.json')
|
|
if os.path.exists(iter_file) and not self.args.override_exist_file:
|
|
return
|
|
|
|
index_resume = -1
|
|
write_mode = 'w'
|
|
if self.args.resume:
|
|
write_mode = 'a'
|
|
if os.path.exists(resume_file):
|
|
shutil.copyfile(resume_file, tmp_file)
|
|
|
|
if os.path.exists(ckpt_state_file):
|
|
with open(ckpt_state_file, 'r', encoding='utf-8') as ckpt_state:
|
|
data = json.load(ckpt_state)
|
|
index_resume = data.get('index', -1)
|
|
logger.info(f'Loaded index_resume: {index_resume}')
|
|
else:
|
|
if os.path.exists(tmp_file):
|
|
os.remove(tmp_file)
|
|
|
|
dataset = self._get_dataset()
|
|
dataset_len = len(dataset)
|
|
total_iters = int(dataset_len // self.args.num_sampling_batch_size)
|
|
|
|
if self.args.num_sampling_batches is None or self.args.num_sampling_batches > total_iters:
|
|
self.args.num_sampling_batches = total_iters
|
|
|
|
with open(tmp_file, write_mode) as f:
|
|
for _index in range(self.args.num_sampling_batches):
|
|
if _index <= index_resume:
|
|
continue
|
|
logger.info(f' Sampling index:{_index}')
|
|
slices = dataset[self.args.num_sampling_batch_size * _index:self.args.num_sampling_batch_size
|
|
* (_index + 1)]
|
|
slices = self.sampler.truncate_input(slices)
|
|
generated = self.sampler.do_sample(slices)
|
|
f.writelines(generated)
|
|
f.flush()
|
|
shutil.copy(tmp_file, resume_file)
|
|
with open(ckpt_state_file, 'w') as ckpt_state:
|
|
json.dump({'index': _index}, ckpt_state)
|
|
|
|
if os.path.exists(iter_file):
|
|
shutil.move(iter_file, iter_file + '.' + str(int(time.time())))
|
|
shutil.move(resume_file, iter_file)
|
|
logger.info(f'Sample file {iter_file} generated.')
|
|
|
|
|
|
def sampling_main(args: Optional[Union[List[str], SamplingArguments]] = None):
|
|
return SwiftSampling(args).main()
|