From ba95f937560aaeffa45b90189653d81200ab4d14 Mon Sep 17 00:00:00 2001
From: Ingmar Schuster <ingmar.schuster@zalando.de>
Date: Mon, 18 May 2020 16:14:46 +0000
Subject: [PATCH] Enabling loading one line at a time

---
 pts/dataset/file_dataset.py | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/pts/dataset/file_dataset.py b/pts/dataset/file_dataset.py
index e48d8bf..cd4ae52 100644
--- a/pts/dataset/file_dataset.py
+++ b/pts/dataset/file_dataset.py
@@ -53,13 +53,19 @@ class JsonLinesFile:
 
     def __init__(self, path: Path, shuffle: bool = True) -> None:
         self.path = path
-        self.shuffle = shuffle
-
+        if shuffle:
+            self.__iter_internal = self.__iter_shuffle__
+        else:
+            self.__iter_internal = self.__iter_inorder__
+    
     def __iter__(self):
+        return self.__iter_internal()
+
+    def __iter_shuffle__(self):
         with open(self.path) as jsonl_file:
             lines = jsonl_file.read().splitlines()
-            if self.shuffle:
-                random.shuffle(lines)
+
+            random.shuffle(lines)
 
             for line_number, raw in enumerate(lines, start=1):
                 span = Span(path=self.path, line=line_number)
@@ -67,6 +73,18 @@ class JsonLinesFile:
                     yield Line(json.loads(raw), span=span)
                 except ValueError:
                     raise Exception(f"Could not read json line {line_number}, {raw}")
+    
+    def __iter_inorder__(self):
+        with open(self.path) as jsonl_file:
+            line_number = 0
+            while True:
+                line_number += 1
+                raw = jsonl_file.readline()
+                span = Span(path=self.path, line=line_number)
+                try:
+                    yield Line(json.loads(raw), span=span)
+                except ValueError:
+                    raise Exception(f"Could not read json line {line_number}, {raw}")
 
     def __len__(self):
         # 1MB