Update generative-proof-of-concept-CPU-preprocessing-in-memory.py

Fix iteration on the Dataset.

Author: david-thrower

generative-proof-of-concept-CPU-preprocessing-in-memory.py

@@ -1330,7 +1330,9 @@ def _expand_next_batch(self):
                 end_idx = min(start_idx + self.sample_expansion_batch_size, len(self.raw_text_samples))
                 collect()
                 if start_idx >= end_idx:
-                    raise StopIteration("No more raw samples to process.")
+                    self.current_index = 0 # raise StopIteration("No more raw samples to process.")
+                    start_idx = 0
+                    end_idx = min(self.sample_expansion_batch_size, len(self.raw_text_samples))
 
                 batch_samples = self.raw_text_samples[start_idx:end_idx]
                 self.current_index = end_idx