72 lines
1.5 KiB
Go
72 lines
1.5 KiB
Go
package common
|
|
|
|
import (
|
|
"strings"
|
|
|
|
"github.com/ollama/ollama/llm"
|
|
)
|
|
|
|
func FindStop(sequence string, stops []string) (bool, string) {
|
|
for _, stop := range stops {
|
|
if strings.Contains(sequence, stop) {
|
|
return true, stop
|
|
}
|
|
}
|
|
|
|
return false, ""
|
|
}
|
|
|
|
func ContainsStopSuffix(sequence string, stops []string) bool {
|
|
for _, stop := range stops {
|
|
for i := 1; i <= len(stop); i++ {
|
|
if strings.HasSuffix(sequence, stop[:i]) {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
// truncateStop removes the provided stop string from pieces,
|
|
// returning the partial pieces with stop removed, including truncating
|
|
// the last piece if required (and signalling if this was the case)
|
|
func TruncateStop(resps []llm.CompletionResponse, stop string) ([]llm.CompletionResponse, bool) {
|
|
var sequence string
|
|
for _, resp := range resps {
|
|
sequence += resp.Content
|
|
}
|
|
|
|
idx := strings.Index(sequence, stop)
|
|
if idx < 0 {
|
|
return resps, false
|
|
}
|
|
|
|
truncated := sequence[:idx]
|
|
if len(truncated) == 0 {
|
|
return nil, true
|
|
}
|
|
|
|
result := make([]llm.CompletionResponse, 0, len(resps))
|
|
|
|
// Track position in truncated sequence
|
|
pos := 0
|
|
truncationHappened := false
|
|
for _, resp := range resps {
|
|
if pos >= len(truncated) {
|
|
break
|
|
}
|
|
|
|
chunk := truncated[pos:min(pos+len(resp.Content), len(truncated))]
|
|
if len(chunk) < len(resp.Content) {
|
|
truncationHappened = true
|
|
}
|
|
if len(chunk) > 0 {
|
|
result = append(result, llm.CompletionResponse{Content: chunk})
|
|
}
|
|
pos += len(resp.Content)
|
|
}
|
|
|
|
return result, truncationHappened
|
|
}
|