| | |
| | import os |
| | import shutil |
| | import time |
| | from typing import List, Union |
| |
|
| | import json |
| |
|
| | from swift.llm import SamplingArguments, SwiftPipeline, load_dataset |
| | from swift.utils import get_logger |
| |
|
| | logger = get_logger() |
| |
|
| |
|
| | class SwiftSampling(SwiftPipeline): |
| | args_class = SamplingArguments |
| | args: args_class |
| |
|
| | def __init__(self, args: Union[List[str], SamplingArguments, None] = 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': |
| | from swift.llm.sampling.vanilla_sampler import VanillaSampler |
| | self.sampler = VanillaSampler(self.args) |
| | elif self.args.sampler_type == 'mcts': |
| | from swift.llm.sampling.mcts import MctsSampler |
| | self.sampler = MctsSampler(self.args) |
| | elif self.args.sampler_type == 'distill': |
| | from swift.llm.sampling.distill_sampler import DistillSampler |
| | 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') 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_per_gpu_batch_size) |
| |
|
| | if self.args.num_sampling_per_gpu_batches is None or self.args.num_sampling_per_gpu_batches > total_iters: |
| | self.args.num_sampling_per_gpu_batches = total_iters |
| |
|
| | with open(tmp_file, write_mode) as f: |
| | for _index in range(self.args.num_sampling_per_gpu_batches): |
| | if _index <= index_resume: |
| | continue |
| | logger.info(f' Sampling index:{_index}') |
| | slices = dataset[self.args.num_sampling_per_gpu_batch_size |
| | * _index:self.args.num_sampling_per_gpu_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: Union[List[str], SamplingArguments, None] = None): |
| | return SwiftSampling(args).main() |
| |
|
