Merge branch 'master' into suqi-localtest

Author: hunterhector

docs/spelling_wordlist.txt

@@ -136,3 +136,4 @@ embedding
 embeddings
 docstrings
 numpy
+jsonpickle

examples/content_rewriter/pipeline.py

@@ -29,9 +29,9 @@ def do_process(input_pack_str: str):
     # Let's assume there is a JSON string for us to use.
     datapack: DataPack = pipeline.process([input_pack_str])
     # You can get the JSON form like this.
-    data_json = datapack.serialize()
+    data_json = datapack.to_string()
     # Let's write it out.
-    with open("generation.txt", "w") as fo:
+    with open("generation.txt", "w", encoding="utf-8") as fo:
         fo.write(data_json)
 
 
@@ -42,6 +42,6 @@ def do_process(input_pack_str: str):
     # You should initialize the model here, so we only do it once.
     pipeline.initialize()
 
-    with open("rewriting_input.json") as fi:
+    with open("rewriting_input.json", encoding="utf-8") as fi:
         test_str = fi.read()
         do_process(test_str)

forte/data/base_pack.py

@@ -13,8 +13,11 @@
 # limitations under the License.
 
 import copy
+import gzip
+import pickle
 import uuid
 from abc import abstractmethod
+from pathlib import Path
 from typing import (
     List,
     Optional,
@@ -162,19 +165,44 @@ def pack_name(self, pack_name: str):
         self._meta.pack_name = pack_name
 
     @classmethod
-    def _deserialize(cls, string: str) -> "PackType":
+    def _deserialize(
+        cls,
+        data_source: Union[Path, str],
+        serialize_method: str = "jsonpickle",
+        zip_pack: bool = False,
+    ) -> "PackType":
         """
         This function should deserialize a Pack from a string. The
-         implementation should decide the specific pack type.
+        implementation should decide the specific pack type.
 
         Args:
-            string: The serialized string to be deserialized.
+            data_source: The data path containing pack data. The content
+              of the data could be string or bytes depending on the method of
+              serialization.
+            serialize_method: The method used to serialize the data, this
+              should be the same as how serialization is done. The current
+              options are "jsonpickle" and "pickle". The default method
+              is "jsonpickle".
+            zip_pack: Boolean value indicating whether the input source is
+              zipped.
 
         Returns:
-            An pack object deserialized from the string.
+            An pack object deserialized from the data.
         """
-        pack = jsonpickle.decode(string)
-        return pack
+        _open = gzip.open if zip_pack else open
+
+        if serialize_method == "jsonpickle":
+            with _open(data_source, mode="rt") as f:  # type: ignore
+                pack = cls.from_string(f.read())
+        else:
+            with _open(data_source, mode="rb") as f:  # type: ignore
+                pack = pickle.load(f)
+
+        return pack  # type: ignore
+
+    @classmethod
+    def from_string(cls, data_content: str) -> "BasePack":
+        return jsonpickle.decode(data_content)
 
     @abstractmethod
     def delete_entry(self, entry: EntryType):
@@ -238,13 +266,76 @@ def add_all_remaining_entries(self, component: Optional[str] = None):
             self.add_entry(entry, c_)
         self._pending_entries.clear()
 
-    def serialize(self, drop_record: Optional[bool] = False) -> str:
-        r"""Serializes a pack to a string."""
+    def to_string(
+        self,
+        drop_record: Optional[bool] = False,
+        json_method: str = "jsonpickle",
+        indent: Optional[int] = None,
+    ) -> str:
+        """
+        Return the string representation (json encoded) of this method.
+
+        Args:
+            drop_record: Whether to drop the creation records, default is False.
+            json_method: What method is used to convert data pack to json.
+              Only supports `json_pickle` for now. Default value is
+              `json_pickle`.
+            indent: The indent used for json string.
+
+        Returns: String representation of the data pack.
+        """
+        if drop_record:
+            self._creation_records.clear()
+            self._field_records.clear()
+        if json_method == "jsonpickle":
+            return jsonpickle.encode(self, unpicklable=True, indent=indent)
+        else:
+            raise ValueError(f"Unsupported JSON method {json_method}.")
+
+    def serialize(
+        self,
+        output_path: Union[str, Path],
+        zip_pack: bool = False,
+        drop_record: bool = False,
+        serialize_method: str = "jsonpickle",
+        indent: Optional[int] = None,
+    ):
+        r"""
+        Serializes the data pack to the provided path. The output of this
+        function depends on the serialization method chosen.
+
+        Args:
+            output_path: The path to write data to.
+            zip_pack: Whether to compress the result with `gzip`.
+            drop_record: Whether to drop the creation records, default is False.
+            serialize_method: The method used to serialize the data. Currently
+              supports "jsonpickle" (outputs str) and Python's built-in
+              "pickle" (outputs bytes).
+            indent: Whether to indent the file if written as JSON.
+
+        Returns: Results of serialization.
+        """
+        if zip_pack:
+            _open = gzip.open
+        else:
+            _open = open  # type:ignore
+
         if drop_record:
             self._creation_records.clear()
             self._field_records.clear()
 
-        return jsonpickle.encode(self, unpicklable=True)
+        if serialize_method == "pickle":
+            with _open(output_path, mode="wb") as pickle_out:
+                pickle.dump(self, pickle_out)  # type:ignore
+        elif serialize_method == "jsonpickle":
+            with _open(output_path, mode="wt", encoding="utf-8") as json_out:
+                json_out.write(
+                    self.to_string(drop_record, "jsonpickle", indent=indent)
+                )
+        else:
+            raise NotImplementedError(
+                f"Unsupported serialization method {serialize_method}"
+            )
 
     def view(self):
         return copy.deepcopy(self)
@@ -457,22 +548,6 @@ def get_ids_from(self, components: List[str]) -> Set[int]:
             valid_component_id |= self.get_ids_by_creator(component)
         return valid_component_id
 
-    def get_ids_by_type_subtype(self, entry_type: Type[EntryType]) -> Set[int]:
-        r"""Look up the type_index with key ``entry_type``.
-
-        Args:
-            entry_type: The type of the entry you are looking for.
-
-        Returns:
-             A set of entry ids. The entries are instances of `entry_type` (
-             and also includes instances of the subclasses of `entry_type`).
-        """
-        subclass_index: Set[int] = set()
-        for index_key, index_val in self._index.iter_type_index():
-            if issubclass(index_key, entry_type):
-                subclass_index.update(index_val)
-        return subclass_index
-
     def _expand_to_sub_types(self, entry_type: Type[EntryType]) -> Set[Type]:
         """
         Return all the types and the sub types that inherit from the provided
@@ -511,7 +586,7 @@ def get_entries_of(
             for tid in self._index.query_by_type(entry_type):
                 yield self.get_entry(tid)
         else:
-            for tid in self.get_ids_by_type_subtype(entry_type):
+            for tid in self._index.query_by_type_subtype(entry_type):
                 yield self.get_entry(tid)
 
     @classmethod

forte/data/base_reader.py

@@ -95,13 +95,16 @@ def default_configs(cls):
         values. Used to replace the missing values of input `configs`
         during pipeline construction.
 
-        .. code-block:: python
+        Here:
+          - zip_pack (bool): whether to zip the results. The default value is
+             False.
+
+          - serialize_method: The method used to serialize the data. Current
+              available options are "jsonpickle" and "pickle". Default is
+              "jsonpickle".
 
-            {
-                "name": "reader"
-            }
         """
-        return {"name": "reader"}
+        return {"zip_pack": False, "serialize_method": "jsonpickle"}
 
     @property
     def pack_type(self):
@@ -323,11 +326,19 @@ def cache_data(self, collection: Any, pack: PackType, append: bool):
 
         logger.info("Caching pack to %s", cache_filename)
         if append:
-            with open(cache_filename, "a", encoding="utf-8") as cache:
-                cache.write(pack.serialize() + "\n")
+            with open(
+                cache_filename,
+                "a",
+                encoding="utf-8",
+            ) as cache:
+                cache.write(pack.to_string() + "\n")
         else:
-            with open(cache_filename, "w", encoding="utf-8") as cache:
-                cache.write(pack.serialize() + "\n")
+            with open(
+                cache_filename,
+                "w",
+                encoding="utf-8",
+            ) as cache:
+                cache.write(pack.to_string() + "\n")
 
     def read_from_cache(
         self, cache_filename: Union[Path, str]
@@ -343,7 +354,7 @@ def read_from_cache(
         logger.info("reading from cache file %s", cache_filename)
         with open(cache_filename, "r", encoding="utf-8") as cache_file:
             for line in cache_file:
-                pack = DataPack.deserialize(line.strip())
+                pack = DataPack.from_string(line.strip())
                 if not isinstance(pack, self.pack_type):
                     raise TypeError(
                         f"Pack deserialized from {cache_filename} "

forte/data/data_pack.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+from pathlib import Path
 from typing import (
     Dict,
     Iterable,
@@ -504,20 +505,30 @@ def get_original_index(
         return Span(orig_begin, orig_end)
 
     @classmethod
-    def deserialize(cls, data_pack_string: str) -> "DataPack":
+    def deserialize(
+        cls,
+        data_source: Union[Path, str],
+        serialize_method: str = "jsonpickle",
+        zip_pack: bool = False,
+    ) -> "DataPack":
         """
         Deserialize a Data Pack from a string. This internally calls the
         internal :meth:`~forte.data.base_pack.BasePack._deserialize` function
         from :class:`~forte.data.base_pack.BasePack`.
 
         Args:
-            data_pack_string: The serialized string of a data pack to be
-              deserialized.
+            data_source: The path storing data source.
+            serialize_method: The method used to serialize the data, this
+              should be the same as how serialization is done. The current
+              options are "jsonpickle" and "pickle". The default method
+              is "jsonpickle".
+            zip_pack: Boolean value indicating whether the input source is
+              zipped.
 
         Returns:
             An data pack object deserialized from the string.
         """
-        return cls._deserialize(data_pack_string)
+        return cls._deserialize(data_source, serialize_method, zip_pack)
 
     def _add_entry(self, entry: EntryType) -> EntryType:
         r"""Force add an :class:`~forte.data.ontology.core.Entry` object to the
@@ -586,25 +597,22 @@ def __add_entry_with_check(
                 f"should be an instance of Annotation, Link, Group of Generics."
             )
 
-        # TODO: duplicate is ill-defined.
-        add_new = allow_duplicate or (entry not in target)
-
-        if add_new:
-            target.add(entry)
-
-            # update the data pack index if needed
-            self._index.update_basic_index([entry])
-            if self._index.link_index_on and isinstance(entry, Link):
-                self._index.update_link_index([entry])
-            if self._index.group_index_on and isinstance(entry, Group):
-                self._index.update_group_index([entry])
-            self._index.deactivate_coverage_index()
-
-            self._pending_entries.pop(entry.tid)
-
-            return entry
-        else:
-            return target[target.index(entry)]
+        if not allow_duplicate:
+            index = target.index(entry)
+            if index < 0:
+                # Return the existing entry if duplicate is not allowed.
+                return target[index]
+
+        target.add(entry)
+        # update the data pack index if needed
+        self._index.update_basic_index([entry])
+        if self._index.link_index_on and isinstance(entry, Link):
+            self._index.update_link_index([entry])
+        if self._index.group_index_on and isinstance(entry, Group):
+            self._index.update_group_index([entry])
+        self._index.deactivate_coverage_index()
+        self._pending_entries.pop(entry.tid)
+        return entry
 
     def delete_entry(self, entry: EntryType):
         r"""Delete an :class:`~forte.data.ontology.core.Entry` object from the
@@ -762,7 +770,7 @@ def get_data(
             context_type_, context_args
         )
 
-        valid_context_ids: Set[int] = self.get_ids_by_type_subtype(
+        valid_context_ids: Set[int] = self._index.query_by_type_subtype(
             context_type_
         )
         if context_components: