121 lines
2.4 KiB
Go
121 lines
2.4 KiB
Go
package pipeline
|
|
|
|
import (
|
|
"regexp"
|
|
"strings"
|
|
"time"
|
|
"unicode/utf8"
|
|
)
|
|
|
|
type TextChunker struct {
|
|
MinChars int
|
|
MaxChars int
|
|
StartAfter int
|
|
FirstTimeout time.Duration
|
|
Punctuation bool
|
|
buf strings.Builder
|
|
firstAt time.Time
|
|
}
|
|
|
|
func NewTextChunker(minChars, maxChars, startAfter int, firstTimeout time.Duration, punctuation bool) *TextChunker {
|
|
if minChars <= 0 {
|
|
minChars = 30
|
|
}
|
|
if maxChars < minChars {
|
|
maxChars = 160
|
|
}
|
|
if startAfter <= 0 {
|
|
startAfter = 60
|
|
}
|
|
if firstTimeout <= 0 {
|
|
firstTimeout = 1200 * time.Millisecond
|
|
}
|
|
return &TextChunker{MinChars: minChars, MaxChars: maxChars, StartAfter: startAfter, FirstTimeout: firstTimeout, Punctuation: punctuation}
|
|
}
|
|
|
|
func (c *TextChunker) Add(delta string, final bool) []string {
|
|
delta = normalizeForSpeech(delta)
|
|
if delta != "" {
|
|
if c.buf.Len() == 0 {
|
|
c.firstAt = time.Now()
|
|
}
|
|
c.buf.WriteString(delta)
|
|
}
|
|
var out []string
|
|
for {
|
|
s := strings.TrimSpace(c.buf.String())
|
|
if s == "" {
|
|
c.buf.Reset()
|
|
return out
|
|
}
|
|
emitAt := c.emitIndex(s, final)
|
|
if emitAt <= 0 {
|
|
return out
|
|
}
|
|
chunk := strings.TrimSpace(s[:emitAt])
|
|
out = append(out, chunk)
|
|
rest := strings.TrimSpace(s[emitAt:])
|
|
c.buf.Reset()
|
|
c.buf.WriteString(rest)
|
|
if rest == "" {
|
|
return out
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *TextChunker) Flush() []string { return c.Add("", true) }
|
|
|
|
func (c *TextChunker) emitIndex(s string, final bool) int {
|
|
if final {
|
|
return len(s)
|
|
}
|
|
if c.Punctuation && utf8.RuneCountInString(s) >= c.MinChars {
|
|
if idx := lastSentenceBoundary(s, c.MaxChars); idx > 0 {
|
|
return idx
|
|
}
|
|
}
|
|
if utf8.RuneCountInString(s) >= c.MaxChars {
|
|
return byteIndexByRunes(s, c.MaxChars)
|
|
}
|
|
if utf8.RuneCountInString(s) >= c.StartAfter && time.Since(c.firstAt) >= c.FirstTimeout {
|
|
return len(s)
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func lastSentenceBoundary(s string, maxRunes int) int {
|
|
limit := byteIndexByRunes(s, maxRunes)
|
|
if limit <= 0 || limit > len(s) {
|
|
limit = len(s)
|
|
}
|
|
last := -1
|
|
for i, r := range s[:limit] {
|
|
if r == '.' || r == '?' || r == '!' || r == '…' {
|
|
last = i + len(string(r))
|
|
}
|
|
}
|
|
return last
|
|
}
|
|
|
|
func byteIndexByRunes(s string, n int) int {
|
|
if n <= 0 {
|
|
return 0
|
|
}
|
|
i := 0
|
|
for pos := range s {
|
|
if i == n {
|
|
return pos
|
|
}
|
|
i++
|
|
}
|
|
return len(s)
|
|
}
|
|
|
|
var markdownPattern = regexp.MustCompile(`[*_` + "`" + `#>\[\]]`)
|
|
|
|
func normalizeForSpeech(s string) string {
|
|
s = markdownPattern.ReplaceAllString(s, "")
|
|
s = strings.ReplaceAll(s, "\n", " ")
|
|
return s
|
|
}
|