Add data type property to rows

Author: Enkidu93

machine/corpora/__init__.py

@@ -4,6 +4,7 @@
 from .alignment_row import AlignmentRow
 from .corpora_utils import batch
 from .corpus import Corpus
+from .data_type import DataType
 from .dbl_bundle_text_corpus import DblBundleTextCorpus
 from .dictionary_alignment_corpus import DictionaryAlignmentCorpus
 from .dictionary_text_corpus import DictionaryTextCorpus
@@ -102,6 +103,7 @@
     "batch",
     "Corpus",
     "create_versification_ref_corpus",
+    "DataType",
     "DblBundleTextCorpus",
     "DictionaryAlignmentCorpus",
     "DictionaryTextCorpus",

machine/corpora/corpora_utils.py

@@ -49,14 +49,14 @@ def get_split_indices(
     return set(rand.sample(range(corpus_size), min(split_size, corpus_size)))
 
 
-def get_files(file_patterns: Iterable[str]) -> Iterable[Tuple[str, str]]:
+def get_files(file_patterns: Iterable[str]) -> Iterable[Tuple[str, str, int]]:
     file_patterns = list(file_patterns)
     if len(file_patterns) == 1 and os.path.isfile(file_patterns[0]):
-        yield ("*all*", file_patterns[0])
+        yield ("*all*", file_patterns[0], 0)
     else:
         for i, file_pattern in enumerate(file_patterns):
             if os.path.isfile(file_pattern):
-                yield (str(i), file_pattern)
+                yield (str(i), file_pattern, i)
                 continue
 
             if "*" not in file_pattern and "?" not in file_pattern and not os.path.exists(file_pattern):
@@ -89,7 +89,7 @@ def get_files(file_patterns: Iterable[str]) -> Iterable[Tuple[str, str]]:
                         updated_id += group
                     if len(updated_id) > 0:
                         id = updated_id
-                yield (id, filename)
+                yield (id, filename, i)
 
 
 def gen(iterable: Iterable[T] = []) -> Generator[T, None, None]:

machine/corpora/data_type.py

@@ -0,0 +1,9 @@
+from enum import Enum, auto
+
+
+class DataType(Enum):  # TODO what options to include? Does a verse=SENTENCE for our purposes?
+    GLOSS = auto()
+    PHRASE = auto()
+    SENTENCE = auto()
+    PASSAGE = auto()
+    DOCUMENT = auto()

machine/corpora/memory_text.py

@@ -1,14 +1,18 @@
 from typing import Generator, Iterable
 
 from .corpora_utils import gen
+from .data_type import DataType
 from .text import Text
 from .text_row import TextRow
 
 
 class MemoryText(Text):
-    def __init__(self, id: str, rows: Iterable[TextRow] = []) -> None:
+    def __init__(self, id: str, rows: Iterable[TextRow] = [], data_type: DataType = DataType.SENTENCE) -> None:
         self._id = id
         self._rows = list(rows)
+        if any([r.data_type != data_type for r in self._rows]):
+            raise ValueError(f"{type(data_type)} of rows must match text {type(data_type)} {data_type}")
+        self._data_type = data_type
 
     @property
     def id(self) -> str:
@@ -18,5 +22,9 @@ def id(self) -> str:
     def sort_key(self) -> str:
         return self._id
 
+    @property
+    def data_type(self) -> DataType:
+        return self._data_type
+
     def _get_rows(self) -> Generator[TextRow, None, None]:
         return gen(self._rows)

machine/corpora/n_parallel_text_corpus.py

@@ -2,6 +2,7 @@
 from typing import Any, Callable, Iterable, List, Optional, Sequence, Set, cast
 
 from ..scripture.verse_ref import Versification
+from .data_type import DataType
 from .n_parallel_text_corpus_base import NParallelTextCorpusBase
 from .n_parallel_text_row import NParallelTextRow
 from .scripture_ref import ScriptureRef
@@ -14,6 +15,7 @@ class _RangeRow:
     refs: List[Any]
     segment: List[str]
     is_sentence_start: bool = False
+    data_type: DataType = DataType.SENTENCE
 
     @property
     def is_in_range(self):
@@ -36,6 +38,7 @@ def __init__(self, n: int):
         self.text_id = ""
         self.versifications: Optional[List[Versification]] = None
         self.row_ref_comparer = None
+        self.data_type = DataType.SENTENCE
 
     @property
     def is_in_range(self) -> bool:
@@ -44,6 +47,7 @@ def is_in_range(self) -> bool:
     def add_text_row(self, row: TextRow, index: int):
         self.text_id = row.text_id
         self.rows[index].refs.append(row.ref)
+        self.rows[index].data_type = row.data_type
         if self.rows[index].is_empty:
             self.rows[index].is_sentence_start = row.is_sentence_start
         self.rows[index].segment.extend(row.segment)
@@ -53,6 +57,7 @@ def create_row(self) -> NParallelTextRow:
         reference_refs: List[Any] = [r.refs[0] if len(r.refs) > 0 else None for r in self.rows if len(r.refs) > 0]
         for i in range(len(self.rows)):
             row = self.rows[i]
+            self.data_type = row.data_type
 
             if (
                 self.versifications is not None
@@ -62,7 +67,7 @@ def create_row(self) -> NParallelTextRow:
                 refs[i] = [cast(ScriptureRef, r).change_versification(self.versifications[i]) for r in reference_refs]
             else:
                 refs[i] = row.refs.copy()
-        n_parallel_text_row = NParallelTextRow(self.text_id, refs)
+        n_parallel_text_row = NParallelTextRow(self.text_id, refs, self.data_type)
         n_parallel_text_row.n_segments = [r.segment.copy() for r in self.rows]
         n_parallel_text_row.n_flags = [
             TextRowFlags.SENTENCE_START if r.is_sentence_start else TextRowFlags.NONE for r in self.rows
@@ -288,6 +293,7 @@ def _create_rows(
             yield range_info.create_row()
 
         default_refs = [[r.ref for r in rows if r is not None][0]]
+        data_type = DataType.SENTENCE
 
         text_id: Optional[str] = None
         refs: List[List[Any]] = []
@@ -298,6 +304,7 @@ def _create_rows(
         for i in range(len(rows)):
             row = rows[i]
             if row is not None:
+                data_type = row.data_type
                 text_id = text_id or row.text_id
                 if self.corpora[i].is_scripture:
                     refs[i] = self._correct_versification([row.ref] if row.ref is None else default_refs, i)
@@ -314,7 +321,7 @@ def _create_rows(
                 )
         refs = [r or default_refs for r in refs]
 
-        new_row = NParallelTextRow(cast(str, text_id), refs)
+        new_row = NParallelTextRow(cast(str, text_id), refs, data_type)
         new_row.n_segments = [r.segment if r is not None else [] for r in rows]
         new_row.n_flags = flags
         yield new_row

machine/corpora/n_parallel_text_row.py

@@ -1,17 +1,19 @@
 from typing import Any, Sequence
 
+from .data_type import DataType
 from .text_row import TextRowFlags
 
 
 class NParallelTextRow:
-    def __init__(self, text_id: str, n_refs: Sequence[Sequence[Any]]):
+    def __init__(self, text_id: str, n_refs: Sequence[Sequence[Any]], data_type: DataType = DataType.SENTENCE):
         if len([n_ref for n_ref in n_refs if n_ref is not None and len(n_ref) > 0]) == 0:
             raise ValueError(f"Refs must be provided but n_refs={n_refs}")
         self._text_id = text_id
         self._n_refs = n_refs
         self._n = len(n_refs)
         self.n_segments: Sequence[Sequence[str]] = [[] for _ in range(0, self._n)]
         self.n_flags: Sequence[TextRowFlags] = [TextRowFlags.SENTENCE_START for _ in range(0, self._n)]
+        self._data_type = data_type
 
     @property
     def text_id(self) -> str:
@@ -21,6 +23,10 @@ def text_id(self) -> str:
     def ref(self) -> Any:
         return self._n_refs[0][0]
 
+    @property
+    def data_type(self) -> DataType:
+        return self._data_type
+
     @property
     def n_refs(self) -> Sequence[Sequence[Any]]:
         return self._n_refs
@@ -42,6 +48,6 @@ def text(self, i: int) -> str:
         return " ".join(self.n_segments[i])
 
     def invert(self) -> "NParallelTextRow":
-        inverted_row = NParallelTextRow(self._text_id, list(reversed(self._n_refs)))
+        inverted_row = NParallelTextRow(self._text_id, list(reversed(self._n_refs)), data_type=self.data_type)
         inverted_row.n_flags = list(reversed(self.n_flags))
         return inverted_row

machine/corpora/parallel_text_corpus.py

@@ -26,6 +26,7 @@
 from .aligned_word_pair import AlignedWordPair
 from .corpora_utils import get_split_indices
 from .corpus import Corpus
+from .data_type import DataType
 from .parallel_text_row import ParallelTextRow
 from .token_processors import escape_spaces, lowercase, normalize, unescape_spaces
 
@@ -401,10 +402,11 @@ def to_hf_dataset(
         ref_column: Optional[str] = "ref",
         translation_column: str = "translation",
         alignment_column: Optional[str] = "alignment",
+        data_type_column: Optional[str] = "data_type",
     ) -> Dataset:
         try:
             from datasets.arrow_dataset import Dataset
-            from datasets.features.features import Features, FeatureType, Sequence, Value
+            from datasets.features.features import ClassLabel, Features, FeatureType, Sequence, Value
             from datasets.features.translation import Translation
         except ImportError:
             raise RuntimeError("datasets is not installed.")
@@ -416,6 +418,8 @@ def to_hf_dataset(
             features_dict[ref_column] = Sequence(Value("string"))
         if alignment_column is not None:
             features_dict[alignment_column] = Sequence({source_lang: Value("int32"), target_lang: Value("int32")})
+        if data_type_column is not None:
+            features_dict[data_type_column] = ClassLabel(names=[e.name for e in DataType])
         features = Features(features_dict)
 
         def iterable() -> Iterable[dict]:
@@ -426,6 +430,8 @@ def iterable() -> Iterable[dict]:
                         example[text_id_column] = row.text_id
                     if ref_column is not None:
                         example[ref_column] = row.refs
+                    if data_type_column is not None:
+                        example[data_type_column] = row.data_type.name
                     example[translation_column] = {source_lang: row.source_text, target_lang: row.target_text}
                     if alignment_column is not None:
                         src_indices: List[int] = []

machine/corpora/parallel_text_row.py

@@ -3,6 +3,7 @@
 from typing import Any, Collection, Optional, Sequence
 
 from .aligned_word_pair import AlignedWordPair
+from .data_type import DataType
 from .text_row import TextRowFlags
 
 
@@ -17,6 +18,7 @@ def __init__(
         aligned_word_pairs: Optional[Collection[AlignedWordPair]] = None,
         source_flags: TextRowFlags = TextRowFlags.SENTENCE_START,
         target_flags: TextRowFlags = TextRowFlags.SENTENCE_START,
+        data_type: DataType = DataType.SENTENCE,
     ) -> None:
         if not text_id:
             raise ValueError("A text_id must be set.")
@@ -25,6 +27,7 @@ def __init__(
         self._text_id = text_id
         self._source_refs = source_refs
         self._target_refs = target_refs
+        self._data_type = data_type
         self.source_segment = source_segment
         self.target_segment = target_segment
         self.aligned_word_pairs = aligned_word_pairs
@@ -51,6 +54,10 @@ def ref(self) -> Any:
     def refs(self) -> Sequence[Any]:
         return self.target_refs if len(self.source_refs) == 0 else self.source_refs
 
+    @property
+    def data_type(self) -> DataType:
+        return self._data_type
+
     @property
     def is_source_sentence_start(self) -> bool:
         return TextRowFlags.SENTENCE_START in self.source_flags
@@ -107,4 +114,5 @@ def invert(self) -> ParallelTextRow:
             None if self.aligned_word_pairs is None else [wp.invert() for wp in self.aligned_word_pairs],
             self.target_flags,
             self.source_flags,
+            self.data_type,
         )

machine/corpora/paratext_backup_terms_corpus.py

@@ -2,6 +2,7 @@
 from zipfile import ZipFile
 
 from ..utils.typeshed import StrPath
+from .data_type import DataType
 from .dictionary_text_corpus import DictionaryTextCorpus
 from .key_term import KeyTerm
 from .memory_text import MemoryText
@@ -25,5 +26,12 @@ def __init__(self, filename: StrPath, term_categories: Sequence[str], use_term_g
                 f"{settings.biblical_terms_file_name}"
             )
 
-            text = MemoryText(text_id, [TextRow(text_id, key_term.id, key_term.renderings) for key_term in key_terms])
+            text = MemoryText(
+                text_id,
+                [
+                    TextRow(text_id, key_term.id, key_term.renderings, data_type=DataType.GLOSS)
+                    for key_term in key_terms
+                ],
+                data_type=DataType.GLOSS,
+            )
             self._add_text(text)

machine/corpora/scripture_text.py

@@ -4,14 +4,15 @@
 from ..scripture.verse_ref import VerseRef, Versification
 from ..utils.context_managed_generator import ContextManagedGenerator
 from .corpora_utils import gen, get_scripture_text_sort_key
+from .data_type import DataType
 from .scripture_ref import ScriptureElement, ScriptureRef
 from .text_base import TextBase
 from .text_row import TextRow, TextRowFlags
 
 
 class ScriptureText(TextBase):
     def __init__(self, id: str, versification: Optional[Versification] = None) -> None:
-        super().__init__(id, get_scripture_text_sort_key(id))
+        super().__init__(id, get_scripture_text_sort_key(id), data_type=DataType.SENTENCE)
         self._versification = ENGLISH_VERSIFICATION if versification is None else versification
 
     @property