mirror of
https://github.com/mudler/LocalAI.git
synced 2024-12-29 17:08:52 +00:00
cf747bcdec
* feat: extract output with regexes from LLMs This changset adds `extract_regex` to the LLM config. It is a list of regexes that can match output and will be used to re extract text from the LLM output. This is particularly useful for LLMs which outputs final results into tags. Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * Add tests, enhance output in case of configuration error Signed-off-by: Ettore Di Giacinto <mudler@localai.io> --------- Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
110 lines
2.7 KiB
Go
110 lines
2.7 KiB
Go
package backend_test
|
|
|
|
import (
|
|
. "github.com/mudler/LocalAI/core/backend"
|
|
"github.com/mudler/LocalAI/core/config"
|
|
"github.com/mudler/LocalAI/core/schema"
|
|
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
)
|
|
|
|
var _ = Describe("LLM tests", func() {
|
|
Context("Finetune LLM output", func() {
|
|
var (
|
|
testConfig config.BackendConfig
|
|
input string
|
|
prediction string
|
|
result string
|
|
)
|
|
|
|
BeforeEach(func() {
|
|
testConfig = config.BackendConfig{
|
|
PredictionOptions: schema.PredictionOptions{
|
|
Echo: false,
|
|
},
|
|
LLMConfig: config.LLMConfig{
|
|
Cutstrings: []string{`<.*?>`}, // Example regex for removing XML tags
|
|
ExtractRegex: []string{`<result>(.*?)</result>`}, // Example regex to extract from tags
|
|
TrimSpace: []string{" ", "\n"},
|
|
TrimSuffix: []string{".", "!"},
|
|
},
|
|
}
|
|
})
|
|
|
|
Context("when echo is enabled", func() {
|
|
BeforeEach(func() {
|
|
testConfig.Echo = true
|
|
input = "Hello"
|
|
prediction = "World"
|
|
})
|
|
|
|
It("should prepend input to prediction", func() {
|
|
result = Finetune(testConfig, input, prediction)
|
|
Expect(result).To(Equal("HelloWorld"))
|
|
})
|
|
})
|
|
|
|
Context("when echo is disabled", func() {
|
|
BeforeEach(func() {
|
|
testConfig.Echo = false
|
|
input = "Hello"
|
|
prediction = "World"
|
|
})
|
|
|
|
It("should not modify the prediction with input", func() {
|
|
result = Finetune(testConfig, input, prediction)
|
|
Expect(result).To(Equal("World"))
|
|
})
|
|
})
|
|
|
|
Context("when cutstrings regex is applied", func() {
|
|
BeforeEach(func() {
|
|
input = ""
|
|
prediction = "<div>Hello</div> World"
|
|
})
|
|
|
|
It("should remove substrings matching cutstrings regex", func() {
|
|
result = Finetune(testConfig, input, prediction)
|
|
Expect(result).To(Equal("Hello World"))
|
|
})
|
|
})
|
|
|
|
Context("when extract regex is applied", func() {
|
|
BeforeEach(func() {
|
|
input = ""
|
|
prediction = "<response><result>42</result></response>"
|
|
})
|
|
|
|
It("should extract substrings matching the extract regex", func() {
|
|
result = Finetune(testConfig, input, prediction)
|
|
Expect(result).To(Equal("42"))
|
|
})
|
|
})
|
|
|
|
Context("when trimming spaces", func() {
|
|
BeforeEach(func() {
|
|
input = ""
|
|
prediction = " Hello World "
|
|
})
|
|
|
|
It("should trim spaces from the prediction", func() {
|
|
result = Finetune(testConfig, input, prediction)
|
|
Expect(result).To(Equal("Hello World"))
|
|
})
|
|
})
|
|
|
|
Context("when trimming suffixes", func() {
|
|
BeforeEach(func() {
|
|
input = ""
|
|
prediction = "Hello World."
|
|
})
|
|
|
|
It("should trim suffixes from the prediction", func() {
|
|
result = Finetune(testConfig, input, prediction)
|
|
Expect(result).To(Equal("Hello World"))
|
|
})
|
|
})
|
|
})
|
|
})
|