name: Load dataset using huggingface
metadata:
  annotations: 
    author: Alexey Volkov <alexey.volkov@ark-kun.com>
    canonical_location: 'https://raw.githubusercontent.com/Ark-kun/pipeline_components/master/components/datasets/HuggingFace/Load_dataset/component.yaml'
inputs:
- {name: dataset_name, type: String}
outputs:
- {name: dataset_dict, type: HuggingFaceDatasetDict}
- {name: splits, type: JsonArray}
implementation:
  container:
    image: python:3.9
    command:
    - sh
    - -c
    - (PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location
      'datasets==1.6.2' || PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install
      --quiet --no-warn-script-location 'datasets==1.6.2' --user) && "$0" "$@"
    - sh
    - -ec
    - |
      program_path=$(mktemp)
      printf "%s" "$0" > "$program_path"
      python3 -u "$program_path" "$@"
    - |
      def _make_parent_dirs_and_return_path(file_path: str):
          import os
          os.makedirs(os.path.dirname(file_path), exist_ok=True)
          return file_path

      def load_dataset_using_huggingface(
          dataset_name,
          dataset_dict_path,
      ):
          from datasets import load_dataset

          dataset_dict = load_dataset(dataset_name)
          dataset_dict.save_to_disk(dataset_dict_path)
          splits = list(dataset_dict.keys())
          return (splits,)

      def _serialize_json(obj) -> str:
          if isinstance(obj, str):
              return obj
          import json
          def default_serializer(obj):
              if hasattr(obj, 'to_struct'):
                  return obj.to_struct()
              else:
                  raise TypeError("Object of type '%s' is not JSON serializable and does not have .to_struct() method." % obj.__class__.__name__)
          return json.dumps(obj, default=default_serializer, sort_keys=True)

      import argparse
      _parser = argparse.ArgumentParser(prog='Load dataset using huggingface', description='')
      _parser.add_argument("--dataset-name", dest="dataset_name", type=str, required=True, default=argparse.SUPPRESS)
      _parser.add_argument("--dataset-dict", dest="dataset_dict_path", type=_make_parent_dirs_and_return_path, required=True, default=argparse.SUPPRESS)
      _parser.add_argument("----output-paths", dest="_output_paths", type=str, nargs=1)
      _parsed_args = vars(_parser.parse_args())
      _output_files = _parsed_args.pop("_output_paths", [])

      _outputs = load_dataset_using_huggingface(**_parsed_args)

      _output_serializers = [
          _serialize_json,

      ]

      import os
      for idx, output_file in enumerate(_output_files):
          try:
              os.makedirs(os.path.dirname(output_file))
          except OSError:
              pass
          with open(output_file, 'w') as f:
              f.write(_output_serializers[idx](_outputs[idx]))
    args:
    - --dataset-name
    - {inputValue: dataset_name}
    - --dataset-dict
    - {outputPath: dataset_dict}
    - '----output-paths'
    - {outputPath: splits}
