Merge pull request #27 from MadBomber/split-on-sentences

split-on-sentences

lib/baran.rb

@@ -5,6 +5,7 @@
 require_relative "baran/markdown_splitter"
 require_relative "baran/recursive_character_text_splitter"
 require_relative "baran/character_text_splitter"
+require_relative "baran/sentence_text_splitter"
 
 module Baran
   class Error < StandardError; end

lib/baran/sentence_text_splitter.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module Baran
+  class SentenceTextSplitter < TextSplitter
+    def initialize(chunk_size: 1024, chunk_overlap: 64)
+      super(chunk_size: chunk_size, chunk_overlap: chunk_overlap)
+    end
+
+    def splitted(text)
+      # Use a regex to split text based on the specified sentence-ending characters followed by whitespace
+      text.scan(/[^.!?]+[.!?]+(?:\s+)/).map(&:strip)
+    end
+  end
+end

test/test_sentence_text_spliter.rb

@@ -0,0 +1,41 @@
+require 'minitest/unit'
+require 'baran'
+
+MiniTest::Unit.autorun
+
+class TestSentenceTextSplitter < MiniTest::Unit::TestCase
+  def setup
+    @splitter = Baran::SentenceTextSplitter.new(chunk_size: 10, chunk_overlap: 5)
+  end
+
+  def test_chunks
+    story = <<~TEXT
+      Hack and jill
+        went up the hill to fetch
+        a pail of water.  Jack fell
+      down and broke his crown and Jill
+        came tumbling after.
+      The pail went flying!  Was the water spilled?
+        No, the water was splashed on Bo Peep.
+    TEXT
+
+    chunks = @splitter.chunks(story)
+
+    sentences = chunks
+                  .map  { |chunk| 
+                          chunk[:text]
+                            .gsub(/\s+/, ' ')
+                            .strip 
+                        }
+
+    expected  = [
+      "Hack and jill went up the hill to fetch a pail of water.", 
+      "Jack fell down and broke his crown and Jill came tumbling after.", 
+      "The pail went flying!", 
+      "Was the water spilled?", 
+      "No, the water was splashed on Bo Peep."
+    ]
+
+    assert_equal(sentences, expected)
+  end
+end