Avoid duplicate field storage

yzhangcs · yzhangcs · commit 575e27830e93 · 2022-06-02T10:53:38.000+08:00
diff --git a/supar/utils/data.py b/supar/utils/data.py
@@ -34,7 +34,7 @@ class Dataset(torch.utils.data.Dataset):
             Each sentence includes fields obeying the data format defined in ``transform``.
     """
 
-    def __init__(self, transform: Transform, data: Union[List[List], str], **kwargs) -> Dataset:
+    def __init__(self, transform: Transform, data: Union[str, List[List]], **kwargs) -> Dataset:
         super(Dataset, self).__init__()
 
         self.transform = transform
@@ -82,10 +82,10 @@ def build(self, batch_size: int, n_buckets: int = 1, shuffle: bool = False, dist
         # numericalize all fields
         fields = self.transform(self.sentences)
         # NOTE: the final bucket count is roughly equal to n_buckets
-        self.buckets = dict(zip(*kmeans([len(s.transformed[fields[0].name]) for s in self], n_buckets)))
+        self.buckets = dict(zip(*kmeans([len(s.fields[fields[0].name]) for s in self], n_buckets)))
         self.loader = DataLoader(dataset=self,
                                  batch_sampler=Sampler(self.buckets, batch_size, shuffle, distributed),
-                                 collate_fn=lambda x: Batch(x))
+                                 collate_fn=lambda x: Batch(self.transform, x))
         return self
 
 
diff --git a/supar/utils/transform.py b/supar/utils/transform.py
@@ -44,7 +44,7 @@ def __call__(self, sentences):
         # numericalize the fields of each sentence
         for sentence in progress_bar(sentences):
             for f in self.flattened_fields:
-                sentence.transformed[f.name] = f.transform([getattr(sentence, f.name)])[0]
+                sentence.fields[f.name] = f.transform([getattr(sentence, f.name)])[0]
         return self.flattened_fields
 
     def __getitem__(self, index):
@@ -322,12 +322,12 @@ def istree(cls, sequence: List[int], proj: bool = False, multiroot: bool = False
 
     def load(
         self,
-        data: Union[List[List], str],
+        data: Union[str, List[List]],
         lang: Optional[str] = None,
         proj: bool = False,
         max_len: Optional[int] = None,
         **kwargs
-    ) -> List['CoNLLSentence']:
+    ) -> List[CoNLLSentence]:
         r"""
         Loads the data in CoNLL-X format.
         Also supports for loading data from CoNLL-U file with comments and non-integer IDs.
@@ -622,11 +622,11 @@ def track(node):
 
     def load(
         self,
-        data: Union[List[List], str],
+        data: Union[str, List[List]],
         lang: Optional[str] = None,
         max_len: Optional[int] = None,
         **kwargs
-    ) -> List['TreeSentence']:
+    ) -> List[TreeSentence]:
         r"""
         Args:
             data (list[list] or str):
@@ -665,24 +665,22 @@ def load(
 
 class Batch(object):
 
-    def __init__(self, sentences):
+    def __init__(self, transform, sentences):
         self.sentences = sentences
-        self.transformed = {f.name: f.compose([s.transformed[f.name] for s in sentences])
-                            for f in sentences[0].transform.flattened_fields}
-        self.fields = list(self.transformed.keys())
+        self.fields = {f.name: f.compose([s.fields[f.name] for s in sentences]) for f in transform.flattened_fields}
+        self.names = list(self.fields.keys())
 
     def __repr__(self):
-        s = ', '.join([f"{name}" for name in self.fields])
-        return f"{self.__class__.__name__}({s})"
+        return f'{self.__class__.__name__}({", ".join([f"{name}" for name in self.names])})'
 
     def __getitem__(self, index):
-        return self.transformed[self.fields[index]]
+        return self.fields[self.names[index]]
 
     def __getattr__(self, name):
         if name in self.__dict__:
             return self.__dict__[name]
-        if name in self.transformed:
-            return self.transformed[name]
+        if name in self.fields:
+            return self.fields[name]
         if hasattr(self.sentences[0], name):
             return [getattr(s, name) for s in self.sentences]
         raise AttributeError
@@ -691,13 +689,11 @@ def __getattr__(self, name):
 class Sentence(object):
 
     def __init__(self, transform):
-        self.transform = transform
-
         # mapping from each nested field to their proper position
         self.maps = dict()
         # names of each field
         self.keys = set()
-        for i, field in enumerate(self.transform):
+        for i, field in enumerate(transform):
             if not isinstance(field, Iterable):
                 field = [field]
             for f in field:
@@ -706,7 +702,7 @@ def __init__(self, transform):
                     self.keys.add(f.name)
         # original values and numericalized values of each position
         self.values = []
-        self.transformed = {key: None for key in self.keys}
+        self.fields = {key: None for key in self.keys}
 
     def __contains__(self, key):
         return key in self.keys
