|
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"): |
|
|
|
|
|
data_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))[0] |
|
else: |
|
data_args = parser.parse_args_into_dataclasses()[0] |
|
|
|
|
|
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"): |
|
|
|
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"] |
|
|
|
conditions = [] |
|
conditions += [filter_by_item(ner, 2)] |
|
conditions += [filter_by_length(title, 4)] |
|
conditions += [filter_by_item(ingredients, 2)] |
|
conditions += [filter_by_item(steps, 2)] |
|
|
|
|
|
if not all(conditions): |
|
return None |
|
|
|
ner = ", ".join(ner) |
|
ingredients = " <sep> ".join(ingredients) |
|
steps = " <sep> ".join(steps) |
|
|
|
|
|
ner = cleaning(ner, "ner") |
|
title = cleaning(title, "title") |
|
ingredients = cleaning(ingredients, "ingredients") |
|
steps = cleaning(steps, "steps") |
|
|
|
return { |
|
"inputs": ner, |
|
|
|
"targets": f"title: {title} <section> ingredients: {ingredients} <section> 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) |
|
|
|
logger.info(f"Preparation - [before] consists of {len(dataset[subset])} records!") |
|
logger.info(f"Preparation - [after] consists of {len(data_df)} records!") |
|
|
|
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() |
|
|
