New Splitter trail; revamped io_utils (#51)

* new splitter * formatting * revamped io utils * formatting * removed unused import
allenai · Jan 31, 2023 · 96dbedf · 96dbedf
1 parent 2ab965f
commit 96dbedf
Show file tree

Hide file tree

Showing 7 changed files with 375 additions and 123 deletions.
diff --git a/examples/squad.py b/examples/squad.py
@@ -14,7 +14,7 @@
     )
     >> sm.TextToWordsMapper(
         fields=["question", "context", "answers"],
-        splitter="whitespace",
+        splitter="ws",
     )
     >> sm.SingleSequenceStriderMapper(
         field_to_stride=["context"],

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "smashed"
-version = "0.16.0"
+version = "0.17.0"
 description = """\
 SMASHED is a toolkit designed to apply transformations to samples in \
 datasets, such as fields extraction, tokenization, prompting, batching, \

diff --git a/src/smashed/mappers/text.py b/src/smashed/mappers/text.py
@@ -9,6 +9,7 @@
     BlingFireSplitter,
     WhitespacePlusSplitter,
     WhitespaceSplitter,
+    WhitespaceTrailSplitter,
 )
 
 
@@ -69,15 +70,15 @@ class TextToWordsMapper(SingleBaseMapper):
     def __init__(
         self,
         fields: Union[str, Sequence[str]],
-        splitter: Literal[
-            "blingfire", "whitespace", "whitespace_plus"
-        ] = "whitespace_plus",
+        splitter: Literal["blingfire", "ws", "plus", "trail"] = "plus",
     ):
         if splitter == "blingfire":
             self.splitter = BlingFireSplitter()
-        elif splitter == "whitespace_plus":
+        elif splitter == "plus":
             self.splitter = WhitespacePlusSplitter()
-        elif splitter == "whitespace":
+        elif splitter == "trail":
+            self.splitter = WhitespaceTrailSplitter()
+        elif splitter == "ws":
             self.splitter = WhitespaceSplitter()
         else:
             raise ValueError(f"Unknown splitter: {splitter}")