import ast import logging import os import sys from dataclasses import dataclass, field import pandas as pd from sklearn.model_selection import train_test_split from tqdm import tqdm from typing import Dict, List, Optional, Tuple from datasets import load_dataset from transformers import ( HfArgumentParser, ) from data_utils import ( filter_by_lang_regex, filter_by_steps, filter_by_length, filter_by_item, filter_by_num_sents, filter_by_num_tokens, normalizer ) logger = logging.getLogger(__name__) @dataclass class DataArguments: """ Arguments to which dataset we are going to set up. """ output_dir: str = field( default=".", metadata={"help": "The output directory where the config will be written."}, ) dataset_name: str = field( default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."} ) dataset_data_dir: Optional[str] = field( default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} ) cache_dir: Optional[str] = field( default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}, ) def main(): parser = HfArgumentParser([DataArguments]) if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): # If we pass only one argument to the script and it's the path to a json file, # let's parse it to get our arguments. data_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))[0] else: data_args = parser.parse_args_into_dataclasses()[0] # Setup logging logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", handlers=[logging.StreamHandler(sys.stdout)], ) logger.setLevel(logging.INFO) logger.info(f"Preparing the dataset") if data_args.dataset_name is not None: dataset = load_dataset( data_args.dataset_name, data_dir=data_args.dataset_data_dir, cache_dir=data_args.cache_dir ) else: dataset = load_dataset( data_args.dataset_name, cache_dir=data_args.cache_dir ) def cleaning(text, item_type="ner"): # NOTE: DO THE CLEANING LATER text = normalizer(text, do_lowercase=True) return text def recipe_preparation(item_dict): ner = item_dict["ner"] title = item_dict["title"] ingredients = item_dict["ingredients"] steps = item_dict["directions"] condition_1 = filter_by_item(ner, 4) condition_2 = filter_by_length(title, 10) condition_3 = filter_by_item(ingredients, 4) condition_4 = filter_by_item(steps, 2) condition_5 = filter_by_steps(" ".join(steps)) if not all([condition_1, condition_2, condition_3, condition_4, condition_5]): return None ner = ", ".join(ner) ingredients = " ".join(ingredients) steps = " ".join(steps) # Cleaning ner = cleaning(ner, "ner") title = cleaning(title, "title") ingredients = cleaning(ingredients, "ingredients") steps = cleaning(steps, "steps") return { "inputs": ner, "targets": f"title: {title}
ingredients: {ingredients}
directions: {steps}" } if len(dataset.keys()) > 1: for subset in dataset.keys(): data_dict = [] for item in tqdm(dataset[subset], position=0, total=len(dataset[subset])): item = recipe_preparation(item) if item: data_dict.append(item) data_df = pd.DataFrame(data_dict) logger.info(f"Preparation of [{subset}] set consists of {len(data_df)} records!") output_path = os.path.join(data_args.output_dir, f"{subset}.csv") os.makedirs(os.path.dirname(output_path), exist_ok=True) data_df.to_csv(output_path, sep="\t", encoding="utf-8", index=False) logger.info(f"Data saved here {output_path}") else: data_dict = [] subset = list(dataset.keys())[0] for item in tqdm(dataset[subset], position=0, total=len(dataset[subset])): item = recipe_preparation(item) if item: data_dict.append(item) data_df = pd.DataFrame(data_dict) train, test = train_test_split(data_df, test_size=0.05, random_state=101) train = train.reset_index(drop=True) test = test.reset_index(drop=True) logger.info(f"Preparation of [train] set consists of {len(train)} records!") logger.info(f"Preparation of [test] set consists of {len(test)} records!") os.makedirs(data_args.output_dir, exist_ok=True) train.to_csv(os.path.join(data_args.output_dir, "train.csv"), sep="\t", encoding="utf-8", index=False) test.to_csv(os.path.join(data_args.output_dir, "test.csv"), sep="\t", encoding="utf-8", index=False) logger.info(f"Data saved here {data_args.output_dir}") if __name__ == '__main__': main()