|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """Library for processing crawled content and generating tfrecords."""
|
|
|
| import collections
|
| import json
|
| import multiprocessing.pool
|
| import os
|
| import urllib.parse
|
|
|
| import tensorflow as tf, tf_keras
|
|
|
| from official.nlp.data import classifier_data_lib
|
| from official.nlp.tools import tokenization
|
|
|
|
|
| class RawDataProcessor(object):
|
| """Data converter for story examples."""
|
|
|
| def __init__(self,
|
| vocab: str,
|
| do_lower_case: bool,
|
| len_title: int = 15,
|
| len_passage: int = 200,
|
| max_num_articles: int = 5,
|
| include_article_title_in_passage: bool = False,
|
| include_text_snippet_in_example: bool = False):
|
| """Constructs a RawDataProcessor.
|
|
|
| Args:
|
| vocab: Filepath of the BERT vocabulary.
|
| do_lower_case: Whether the vocabulary is uncased or not.
|
| len_title: Maximum number of tokens in story headline.
|
| len_passage: Maximum number of tokens in article passage.
|
| max_num_articles: Maximum number of articles in a story.
|
| include_article_title_in_passage: Whether to include article title in
|
| article passage.
|
| include_text_snippet_in_example: Whether to include text snippet (headline
|
| and article content) in generated tensorflow Examples, for debug usage.
|
| If include_article_title_in_passage=True, title and body will be
|
| separated by [SEP].
|
| """
|
| self.articles = dict()
|
| self.tokenizer = tokenization.FullTokenizer(
|
| vocab, do_lower_case=do_lower_case, split_on_punc=False)
|
| self.len_title = len_title
|
| self.len_passage = len_passage
|
| self.max_num_articles = max_num_articles
|
| self.include_article_title_in_passage = include_article_title_in_passage
|
| self.include_text_snippet_in_example = include_text_snippet_in_example
|
|
|
| self.ex_index = 5
|
|
|
| self.label = 0
|
| self.guid = 0
|
| self.num_generated_examples = 0
|
|
|
| def read_crawled_articles(self, folder_path):
|
| """Reads crawled articles under folder_path."""
|
| for path, _, files in os.walk(folder_path):
|
| for name in files:
|
| if not name.endswith(".json"):
|
| continue
|
| url, article = self._get_article_content_from_json(
|
| os.path.join(path, name))
|
| if not article.text_a:
|
| continue
|
| self.articles[RawDataProcessor.normalize_url(url)] = article
|
| if len(self.articles) % 5000 == 0:
|
| print("Number of articles loaded: %d\r" % len(self.articles), end="")
|
| print()
|
| return len(self.articles)
|
|
|
| def generate_examples(self, input_file, output_files):
|
| """Loads story from input json file and exports examples in output_files."""
|
| writers = []
|
| story_partition = []
|
| for output_file in output_files:
|
| writers.append(tf.io.TFRecordWriter(output_file))
|
| story_partition.append(list())
|
| with tf.io.gfile.GFile(input_file, "r") as story_json_file:
|
| stories = json.load(story_json_file)
|
| writer_index = 0
|
| for story in stories:
|
| articles = []
|
| for url in story["urls"]:
|
| normalized_url = RawDataProcessor.normalize_url(url)
|
| if normalized_url in self.articles:
|
| articles.append(self.articles[normalized_url])
|
| if not articles:
|
| continue
|
| story_partition[writer_index].append((story["label"], articles))
|
| writer_index = (writer_index + 1) % len(writers)
|
| lock = multiprocessing.Lock()
|
| pool = multiprocessing.pool.ThreadPool(len(writers))
|
| data = [(story_partition[i], writers[i], lock) for i in range(len(writers))]
|
| pool.map(self._write_story_partition, data)
|
| return len(stories), self.num_generated_examples
|
|
|
| @classmethod
|
| def normalize_url(cls, url):
|
| """Normalize url for better matching."""
|
| url = urllib.parse.unquote(
|
| urllib.parse.urlsplit(url)._replace(query=None).geturl())
|
| output, part = [], None
|
| for part in url.split("//"):
|
| if part == "http:" or part == "https:":
|
| continue
|
| else:
|
| output.append(part)
|
| return "//".join(output)
|
|
|
| def _get_article_content_from_json(self, file_path):
|
| """Returns (url, InputExample) keeping content extracted from file_path."""
|
| with tf.io.gfile.GFile(file_path, "r") as article_json_file:
|
| article = json.load(article_json_file)
|
| if self.include_article_title_in_passage:
|
| return article["url"], classifier_data_lib.InputExample(
|
| guid=self.guid,
|
| text_a=article["title"],
|
| text_b=article["maintext"],
|
| label=self.label)
|
| else:
|
| return article["url"], classifier_data_lib.InputExample(
|
| guid=self.guid, text_a=article["maintext"], label=self.label)
|
|
|
| def _write_story_partition(self, data):
|
| """Writes stories in a partition into file."""
|
| for (story_headline, articles) in data[0]:
|
| story_example = tf.train.Example(
|
| features=tf.train.Features(
|
| feature=self._get_single_story_features(story_headline,
|
| articles)))
|
| data[1].write(story_example.SerializeToString())
|
| data[2].acquire()
|
| try:
|
| self.num_generated_examples += 1
|
| if self.num_generated_examples % 1000 == 0:
|
| print(
|
| "Number of stories written: %d\r" % self.num_generated_examples,
|
| end="")
|
| finally:
|
| data[2].release()
|
|
|
| def _get_single_story_features(self, story_headline, articles):
|
| """Converts a list of articles to a tensorflow Example."""
|
|
|
| def get_text_snippet(article):
|
| if article.text_b:
|
| return " [SEP] ".join([article.text_a, article.text_b])
|
| else:
|
| return article.text_a
|
|
|
| story_features = collections.OrderedDict()
|
| story_headline_feature = classifier_data_lib.convert_single_example(
|
| ex_index=self.ex_index,
|
| example=classifier_data_lib.InputExample(
|
| guid=self.guid, text_a=story_headline, label=self.label),
|
| label_list=[self.label],
|
| max_seq_length=self.len_title,
|
| tokenizer=self.tokenizer)
|
| if self.include_text_snippet_in_example:
|
| story_headline_feature.label_id = story_headline
|
| self._add_feature_with_suffix(
|
| feature=story_headline_feature,
|
| suffix="a",
|
| story_features=story_features)
|
| for (article_index, article) in enumerate(articles):
|
| if article_index == self.max_num_articles:
|
| break
|
| article_feature = classifier_data_lib.convert_single_example(
|
| ex_index=self.ex_index,
|
| example=article,
|
| label_list=[self.label],
|
| max_seq_length=self.len_passage,
|
| tokenizer=self.tokenizer)
|
| if self.include_text_snippet_in_example:
|
| article_feature.label_id = get_text_snippet(article)
|
| suffix = chr(ord("b") + article_index)
|
| self._add_feature_with_suffix(
|
| feature=article_feature, suffix=suffix, story_features=story_features)
|
|
|
|
|
| for article_index in range(len(articles), self.max_num_articles):
|
| suffix = chr(ord("b") + article_index)
|
| empty_article = classifier_data_lib.InputExample(
|
| guid=self.guid, text_a="", label=self.label)
|
| empty_feature = classifier_data_lib.convert_single_example(
|
| ex_index=self.ex_index,
|
| example=empty_article,
|
| label_list=[self.label],
|
| max_seq_length=self.len_passage,
|
| tokenizer=self.tokenizer)
|
| if self.include_text_snippet_in_example:
|
| empty_feature.label_id = ""
|
| self._add_feature_with_suffix(
|
| feature=empty_feature, suffix=suffix, story_features=story_features)
|
| return story_features
|
|
|
| def _add_feature_with_suffix(self, feature, suffix, story_features):
|
| """Appends suffix to feature names and fills in the corresponding values."""
|
|
|
| def _create_int_feature(values):
|
| return tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
|
|
|
| def _create_string_feature(value):
|
| return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
|
|
|
| story_features["input_ids_%c" % suffix] = _create_int_feature(
|
| feature.input_ids)
|
| story_features["input_mask_%c" % suffix] = _create_int_feature(
|
| feature.input_mask)
|
| story_features["segment_ids_%c" % suffix] = _create_int_feature(
|
| feature.segment_ids)
|
| if self.include_text_snippet_in_example:
|
| story_features["text_snippet_%c" % suffix] = _create_string_feature(
|
| bytes(feature.label_id.encode()))
|
|
|
