Merge pull request #5 from kylepjohnson/main

Update split() to return list of lists
agamm · Jul 11, 2024 · a9f2310 · a9f2310
2 parents 18558be + 272005c
commit a9f2310
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/semantic_split/splitters/SimilarSentenceSplitter.py b/semantic_split/splitters/SimilarSentenceSplitter.py
@@ -11,7 +11,7 @@ def __init__(self, similarity_model, sentence_splitter: Splitter):
         self.model = similarity_model
         self.sentence_splitter = sentence_splitter
 
-    def split(self, text: str, group_max_sentences=5) -> List[str]:
+    def split(self, text: str, group_max_sentences=5) -> List[List[str]]:
         '''
             group_max_sentences: The maximum number of sentences in a group.
         '''
@@ -36,4 +36,4 @@ def split(self, text: str, group_max_sentences=5) -> List[str]:
             else:
                 groups.append([sentences[i]])
 
-        return groups
+        return groups