text: # Copyright 2020 The Forte Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """
  The reader that reads AG News data into Datapacks.
  """
  from typing import Iterator, Tuple
  
  from ftx.onto.ag_news import Description
  from ft.onto.base_ontology import Document, Title
  from forte.data.data_pack import DataPack
  from forte.data.base_reader import PackReader
  
  
  __all__ = [
      "AGNewsReader",
  ]
  
  
  class AGNewsReader(PackReader):
      r""":class:`AGNewsReader` is designed to read in AG News
      text classification dataset.
      The AG's news topic classification dataset is constructed by Xiang Zhang
      (xiang.zhang@nyu.edu) from the AG corpus. It is used as a text
      classification benchmark in the following paper:
      https://arxiv.org/abs/1509.01626
      The dataset can be downloaded from:
      https://github.com/mhjabreel/CharCnn_Keras/tree/master/data/ag_news_csv
  
      The input to this reader is the path to the CSV file.
      """
  
      def _collect(  # type: ignore
          self, csv_file: str
      ) -> Iterator[Tuple[int, str]]:
          r"""Collects from a CSV file path and returns an iterator of AG News
          data. The elements in the iterator correspond to each line
          in the csv file. One line is expected to be parsed as one
          DataPack.
  
          Args:
              csv_file: A CSV file path.
  
          Returns: Iterator of each line in the csv file.
          """
          with open(csv_file, "r", encoding="utf-8") as f:
              for line_id, line in enumerate(f):
                  yield line_id, line
  
      def _cache_key_function(self, line_info: Tuple[int, str]) -> str:
          return str(line_info[0])
  
      def _parse_pack(self, line_info: Tuple[int, str]) -> Iterator[DataPack]:
          line_id, line = line_info
  
          pack = DataPack()
          text: str = ""
          line = line.strip()
          data = line.split(",")
  
          class_id: int = int(data[0].replace('"', ""))
          title: str = data[1]
          description: str = data[2]
  
          text += title
          title_end = len(text)
          text += "\n" + description
          description_start = title_end + 1
          description_end = len(text)
  
          pack.set_text(text, replace_func=self.text_replace_operation)
  
          doc = Document(pack, 0, description_end)
          doc.document_class = [self.configs.class_names[class_id]]
          Title(pack, 0, title_end)
          Description(pack, description_start, description_end)
  
          pack.pack_name = line_id
          yield pack
  
      @classmethod
      def default_configs(cls):
          config: dict = super().default_configs()
  
          config.update(
              {
                  "class_names": {
                      1: "World",
                      2: "Sports",
                      3: "Business",
                      4: "Sci/Tech",
                  }
              }
          )
          return config
  