wehub-resource-sync
48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import os
|
|
import torch
|
|
from typing import List, Optional, Union
|
|
|
|
from swift.arguments import ExportArguments
|
|
from swift.utils import get_logger
|
|
from ..train import SwiftSft
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class ExportCachedDataset(SwiftSft):
|
|
args_class = ExportArguments
|
|
args: args_class
|
|
|
|
def __init__(self, args: Optional[Union[List[str], ExportArguments]] = None) -> None:
|
|
super(SwiftSft, self).__init__(args)
|
|
args = self.args
|
|
self.train_msg = {} # dummy
|
|
template_cls = args.template_meta.template_cls
|
|
if template_cls and template_cls.use_model:
|
|
kwargs = {'return_dummy_model': True}
|
|
else:
|
|
kwargs = {'load_model': False}
|
|
with torch.device('meta'):
|
|
self._prepare_model_tokenizer(**kwargs)
|
|
self._prepare_template()
|
|
self.template.set_mode(args.template_mode)
|
|
|
|
def _post_process_datasets(self, datasets: List) -> List:
|
|
return datasets
|
|
|
|
def main(self):
|
|
train_dataset, val_dataset = self._prepare_dataset()
|
|
train_data_dir = os.path.join(self.args.output_dir, 'train')
|
|
val_data_dir = os.path.join(self.args.output_dir, 'val')
|
|
train_dataset.save_to_disk(train_data_dir)
|
|
if val_dataset is not None:
|
|
val_dataset.save_to_disk(val_data_dir)
|
|
logger.info(f'cached_dataset: `{train_data_dir}`')
|
|
if val_dataset is not None:
|
|
logger.info(f'cached_val_dataset: `{val_data_dir}`')
|
|
|
|
|
|
def export_cached_dataset(args: Optional[Union[List[str], ExportArguments]] = None):
|
|
return ExportCachedDataset(args).main()
|