Added iterators for clusters.

wanmok · Apr 7, 2021 · 804e04a · 804e04a
1 parent 7acba52
commit 804e04a
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/cement/cement_document.py b/cement/cement_document.py
@@ -463,6 +463,18 @@ def iterate_sentences(self) -> Iterable[List[str]]:
         for i, offset in enumerate(self._tokenization_offsets):
             yield self[(0 if i == 0 else self._tokenization_offsets[i - 1]):offset]
 
+    def iterate_entities(self, key: Callable = lambda x: True) -> Iterable[Entity]:
+        return filter(key, self._entity_set.entityList)
+
+    def iterate_situations(self, key: Callable = lambda x: True) -> Iterable[Situation]:
+        return filter(key, self._situation_set.situationList)
+
+    def iterate_events(self) -> Iterable[Situation]:
+        return self.iterate_situations(key=lambda x: x.situationType == 'EVENT')
+
+    def iterate_relations(self) -> Iterable[Situation]:
+        return self.iterate_situations(key=lambda x: x.situationType == 'RELATION')
+
     def iterate_situation_mentions(self, key: Callable = lambda x: True) -> Iterable[SituationMention]:
         return filter(key, self._situation_mention_set.mentionList)