TokenTextSplitter Ref: extract method, use range.

xebia-functional · Jun 22, 2023 · c24450f · c24450f
1 parent 3a7a08d
commit c24450f
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/textsplitters/TokenTextSplitter.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/textsplitters/TokenTextSplitter.kt
@@ -2,6 +2,7 @@ package com.xebia.functional.xef.textsplitters
 
 import com.xebia.functional.tokenizer.Encoding
 import com.xebia.functional.tokenizer.ModelType
+import kotlin.ranges.IntProgression
 
 fun TokenTextSplitter(modelType: ModelType, chunkSize: Int, chunkOverlap: Int): TextSplitter =
   TokenTextSplitterImpl(modelType.encoding, chunkSize, chunkOverlap)
@@ -16,12 +17,14 @@ private class TokenTextSplitterImpl(
     val inputIds = tokenizer.encode(text)
     val stepSize = chunkSize - chunkOverlap
 
-    return inputIds.indices
-      .asSequence()
-      .filter { it % stepSize == 0 }
-      .map { startIdx -> inputIds.subList(startIdx, minOf(startIdx + chunkSize, inputIds.size)) }
-      .map { chunkIds -> tokenizer.decode(chunkIds) }
-      .toList()
+    fun decodeSegment(startIdx: Int): String {
+      val end = minOf(startIdx + chunkSize, inputIds.size)
+      return tokenizer.decode(inputIds.subList(startIdx, end))
+    }
+
+    return IntProgression.fromClosedRange(0, inputIds.size - 1, stepSize)
+        .toList()
+        .map { seg -> decodeSegment(seg) }
   }
 
   override suspend fun splitDocuments(documents: List<String>): List<String> =