ollama/parser/parser_test.go

502 lines
10 KiB
Go
Raw Normal View History

2024-01-05 22:04:31 +00:00
package parser
import (
2024-04-22 22:37:14 +00:00
"bytes"
"fmt"
"io"
2024-01-05 22:04:31 +00:00
"strings"
"testing"
"github.com/stretchr/testify/assert"
)
2024-04-22 22:37:14 +00:00
func TestParser(t *testing.T) {
2024-01-05 22:04:31 +00:00
input := `
FROM model1
ADAPTER adapter1
LICENSE MIT
PARAMETER param1 value1
PARAMETER param2 value2
TEMPLATE template1
`
reader := strings.NewReader(input)
commands, err := Parse(reader)
assert.Nil(t, err)
expectedCommands := []Command{
{Name: "model", Args: "model1"},
{Name: "adapter", Args: "adapter1"},
{Name: "license", Args: "MIT"},
{Name: "param1", Args: "value1"},
{Name: "param2", Args: "value2"},
{Name: "template", Args: "template1"},
}
assert.Equal(t, expectedCommands, commands)
}
2024-04-24 23:12:56 +00:00
func TestParserFrom(t *testing.T) {
var cases = []struct {
input string
expected []Command
err error
}{
{
"FROM foo",
[]Command{{Name: "model", Args: "foo"}},
nil,
},
{
"FROM /path/to/model",
[]Command{{Name: "model", Args: "/path/to/model"}},
nil,
},
{
"FROM /path/to/model/fp16.bin",
[]Command{{Name: "model", Args: "/path/to/model/fp16.bin"}},
nil,
},
{
"FROM llama3:latest",
[]Command{{Name: "model", Args: "llama3:latest"}},
nil,
},
{
"FROM llama3:7b-instruct-q4_K_M",
[]Command{{Name: "model", Args: "llama3:7b-instruct-q4_K_M"}},
nil,
},
{
"", nil, errMissingFrom,
},
{
"PARAMETER param1 value1",
nil,
errMissingFrom,
},
{
"PARAMETER param1 value1\nFROM foo",
[]Command{{Name: "param1", Args: "value1"}, {Name: "model", Args: "foo"}},
nil,
},
}
2024-01-05 22:04:31 +00:00
2024-04-24 23:12:56 +00:00
for _, c := range cases {
t.Run("", func(t *testing.T) {
commands, err := Parse(strings.NewReader(c.input))
assert.ErrorIs(t, err, c.err)
assert.Equal(t, c.expected, commands)
})
}
2024-01-05 22:04:31 +00:00
}
2024-04-22 22:37:14 +00:00
func TestParserParametersMissingValue(t *testing.T) {
2024-01-05 22:04:31 +00:00
input := `
FROM foo
PARAMETER param1
`
reader := strings.NewReader(input)
_, err := Parse(reader)
2024-04-22 22:37:14 +00:00
assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
2024-01-05 22:04:31 +00:00
}
2024-01-25 20:12:36 +00:00
2024-04-27 00:11:47 +00:00
func TestParserBadCommand(t *testing.T) {
input := `
FROM foo
BADCOMMAND param1 value1
`
_, err := Parse(strings.NewReader(input))
assert.ErrorIs(t, err, errInvalidCommand)
}
2024-04-22 22:37:14 +00:00
func TestParserMessages(t *testing.T) {
var cases = []struct {
input string
expected []Command
err error
}{
{
`
FROM foo
MESSAGE system You are a Parser. Always Parse things.
`,
[]Command{
{Name: "model", Args: "foo"},
{Name: "message", Args: "system: You are a Parser. Always Parse things."},
},
nil,
},
{
`
2024-01-25 20:12:36 +00:00
FROM foo
2024-04-25 02:17:26 +00:00
MESSAGE system You are a Parser. Always Parse things.`,
[]Command{
{Name: "model", Args: "foo"},
{Name: "message", Args: "system: You are a Parser. Always Parse things."},
},
nil,
},
{
`
FROM foo
2024-01-25 20:12:36 +00:00
MESSAGE system You are a Parser. Always Parse things.
MESSAGE user Hey there!
MESSAGE assistant Hello, I want to parse all the things!
2024-04-22 22:37:14 +00:00
`,
[]Command{
{Name: "model", Args: "foo"},
{Name: "message", Args: "system: You are a Parser. Always Parse things."},
{Name: "message", Args: "user: Hey there!"},
{Name: "message", Args: "assistant: Hello, I want to parse all the things!"},
},
nil,
},
{
`
FROM foo
MESSAGE system """
You are a multiline Parser. Always Parse things.
"""
`,
[]Command{
{Name: "model", Args: "foo"},
{Name: "message", Args: "system: \nYou are a multiline Parser. Always Parse things.\n"},
},
nil,
},
{
`
FROM foo
MESSAGE badguy I'm a bad guy!
`,
nil,
2024-04-27 00:11:47 +00:00
errInvalidMessageRole,
2024-04-22 22:37:14 +00:00
},
{
`
FROM foo
MESSAGE system
`,
nil,
io.ErrUnexpectedEOF,
},
{
`
FROM foo
MESSAGE system`,
nil,
io.ErrUnexpectedEOF,
},
}
2024-01-25 20:12:36 +00:00
2024-04-22 22:37:14 +00:00
for _, c := range cases {
t.Run("", func(t *testing.T) {
commands, err := Parse(strings.NewReader(c.input))
assert.ErrorIs(t, err, c.err)
assert.Equal(t, c.expected, commands)
})
}
}
2024-01-25 20:12:36 +00:00
2024-04-22 22:37:14 +00:00
func TestParserQuoted(t *testing.T) {
var cases = []struct {
multiline string
expected []Command
err error
}{
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM """
2024-04-22 22:37:14 +00:00
This is a
2024-04-25 02:17:26 +00:00
multiline system.
2024-04-22 22:37:14 +00:00
"""
`,
[]Command{
{Name: "model", Args: "foo"},
2024-04-25 02:17:26 +00:00
{Name: "system", Args: "\nThis is a\nmultiline system.\n"},
2024-04-22 22:37:14 +00:00
},
nil,
},
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM """
2024-04-22 22:37:14 +00:00
This is a
2024-04-25 02:17:26 +00:00
multiline system."""
2024-04-22 22:37:14 +00:00
`,
[]Command{
{Name: "model", Args: "foo"},
2024-04-25 02:17:26 +00:00
{Name: "system", Args: "\nThis is a\nmultiline system."},
2024-04-22 22:37:14 +00:00
},
nil,
},
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM """This is a
multiline system."""
2024-04-22 22:37:14 +00:00
`,
[]Command{
{Name: "model", Args: "foo"},
2024-04-25 02:17:26 +00:00
{Name: "system", Args: "This is a\nmultiline system."},
2024-04-22 22:37:14 +00:00
},
nil,
},
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM """This is a multiline system."""
2024-04-22 22:37:14 +00:00
`,
[]Command{
{Name: "model", Args: "foo"},
2024-04-25 02:17:26 +00:00
{Name: "system", Args: "This is a multiline system."},
2024-04-22 22:37:14 +00:00
},
nil,
},
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM """This is a multiline system.""
2024-04-22 22:37:14 +00:00
`,
nil,
io.ErrUnexpectedEOF,
},
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM "
2024-04-22 22:37:14 +00:00
`,
nil,
io.ErrUnexpectedEOF,
},
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM """
This is a multiline system with "quotes".
2024-04-22 22:37:14 +00:00
"""
`,
[]Command{
{Name: "model", Args: "foo"},
2024-04-25 02:17:26 +00:00
{Name: "system", Args: "\nThis is a multiline system with \"quotes\".\n"},
2024-04-22 22:37:14 +00:00
},
nil,
},
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM """"""
2024-04-22 22:37:14 +00:00
`,
[]Command{
{Name: "model", Args: "foo"},
2024-04-25 02:17:26 +00:00
{Name: "system", Args: ""},
2024-04-22 22:37:14 +00:00
},
nil,
},
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM ""
2024-04-22 22:37:14 +00:00
`,
[]Command{
{Name: "model", Args: "foo"},
2024-04-25 02:17:26 +00:00
{Name: "system", Args: ""},
2024-04-22 22:37:14 +00:00
},
nil,
},
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM "'"
2024-04-22 22:37:14 +00:00
`,
[]Command{
{Name: "model", Args: "foo"},
2024-04-25 02:17:26 +00:00
{Name: "system", Args: "'"},
2024-04-22 22:37:14 +00:00
},
nil,
},
2024-04-24 23:12:56 +00:00
{
`
FROM foo
2024-04-25 02:17:26 +00:00
SYSTEM """''"'""'""'"'''''""'""'"""
2024-04-24 23:12:56 +00:00
`,
[]Command{
{Name: "model", Args: "foo"},
2024-04-25 02:17:26 +00:00
{Name: "system", Args: `''"'""'""'"'''''""'""'`},
},
nil,
},
{
`
FROM foo
TEMPLATE """
{{ .Prompt }}
"""`,
[]Command{
{Name: "model", Args: "foo"},
{Name: "template", Args: "\n{{ .Prompt }}\n"},
2024-04-24 23:12:56 +00:00
},
nil,
},
2024-01-25 20:12:36 +00:00
}
2024-04-22 22:37:14 +00:00
for _, c := range cases {
t.Run("", func(t *testing.T) {
commands, err := Parse(strings.NewReader(c.multiline))
assert.ErrorIs(t, err, c.err)
assert.Equal(t, c.expected, commands)
})
}
2024-01-25 20:12:36 +00:00
}
2024-04-22 22:37:14 +00:00
func TestParserParameters(t *testing.T) {
2024-04-24 23:12:56 +00:00
var cases = map[string]struct {
name, value string
}{
"numa true": {"numa", "true"},
"num_ctx 1": {"num_ctx", "1"},
"num_batch 1": {"num_batch", "1"},
"num_gqa 1": {"num_gqa", "1"},
"num_gpu 1": {"num_gpu", "1"},
"main_gpu 1": {"main_gpu", "1"},
"low_vram true": {"low_vram", "true"},
"f16_kv true": {"f16_kv", "true"},
"logits_all true": {"logits_all", "true"},
"vocab_only true": {"vocab_only", "true"},
"use_mmap true": {"use_mmap", "true"},
"use_mlock true": {"use_mlock", "true"},
"num_thread 1": {"num_thread", "1"},
"num_keep 1": {"num_keep", "1"},
"seed 1": {"seed", "1"},
"num_predict 1": {"num_predict", "1"},
"top_k 1": {"top_k", "1"},
"top_p 1.0": {"top_p", "1.0"},
"tfs_z 1.0": {"tfs_z", "1.0"},
"typical_p 1.0": {"typical_p", "1.0"},
"repeat_last_n 1": {"repeat_last_n", "1"},
"temperature 1.0": {"temperature", "1.0"},
"repeat_penalty 1.0": {"repeat_penalty", "1.0"},
"presence_penalty 1.0": {"presence_penalty", "1.0"},
"frequency_penalty 1.0": {"frequency_penalty", "1.0"},
"mirostat 1": {"mirostat", "1"},
"mirostat_tau 1.0": {"mirostat_tau", "1.0"},
"mirostat_eta 1.0": {"mirostat_eta", "1.0"},
"penalize_newline true": {"penalize_newline", "true"},
"stop ### User:": {"stop", "### User:"},
"stop ### User: ": {"stop", "### User: "},
"stop \"### User:\"": {"stop", "### User:"},
"stop \"### User: \"": {"stop", "### User: "},
"stop \"\"\"### User:\"\"\"": {"stop", "### User:"},
"stop \"\"\"### User:\n\"\"\"": {"stop", "### User:\n"},
"stop <|endoftext|>": {"stop", "<|endoftext|>"},
"stop <|eot_id|>": {"stop", "<|eot_id|>"},
"stop </s>": {"stop", "</s>"},
2024-04-22 22:37:14 +00:00
}
2024-01-25 20:12:36 +00:00
2024-04-24 23:12:56 +00:00
for k, v := range cases {
t.Run(k, func(t *testing.T) {
2024-04-22 22:37:14 +00:00
var b bytes.Buffer
fmt.Fprintln(&b, "FROM foo")
2024-04-24 23:12:56 +00:00
fmt.Fprintln(&b, "PARAMETER", k)
commands, err := Parse(&b)
2024-04-22 22:37:14 +00:00
assert.Nil(t, err)
2024-04-24 23:12:56 +00:00
assert.Equal(t, []Command{
{Name: "model", Args: "foo"},
{Name: v.name, Args: v.value},
}, commands)
2024-04-22 22:37:14 +00:00
})
}
}
func TestParserComments(t *testing.T) {
var cases = []struct {
input string
expected []Command
}{
{
`
# comment
2024-01-25 20:12:36 +00:00
FROM foo
2024-04-22 22:37:14 +00:00
`,
[]Command{
{Name: "model", Args: "foo"},
},
},
}
2024-01-25 20:12:36 +00:00
2024-04-22 22:37:14 +00:00
for _, c := range cases {
t.Run("", func(t *testing.T) {
commands, err := Parse(strings.NewReader(c.input))
assert.Nil(t, err)
assert.Equal(t, c.expected, commands)
})
}
2024-01-25 20:12:36 +00:00
}
2024-04-25 01:49:14 +00:00
func TestParseFormatParse(t *testing.T) {
var cases = []string{
`
FROM foo
ADAPTER adapter1
LICENSE MIT
PARAMETER param1 value1
PARAMETER param2 value2
TEMPLATE template1
MESSAGE system You are a Parser. Always Parse things.
MESSAGE user Hey there!
MESSAGE assistant Hello, I want to parse all the things!
`,
`
FROM foo
ADAPTER adapter1
LICENSE MIT
PARAMETER param1 value1
PARAMETER param2 value2
TEMPLATE template1
MESSAGE system """
You are a store greeter. Always responsed with "Hello!".
"""
MESSAGE user Hey there!
MESSAGE assistant Hello, I want to parse all the things!
`,
`
FROM foo
ADAPTER adapter1
LICENSE """
Very long and boring legal text.
Blah blah blah.
"Oh look, a quote!"
"""
PARAMETER param1 value1
PARAMETER param2 value2
TEMPLATE template1
MESSAGE system """
You are a store greeter. Always responsed with "Hello!".
"""
MESSAGE user Hey there!
MESSAGE assistant Hello, I want to parse all the things!
`,
}
for _, c := range cases {
t.Run("", func(t *testing.T) {
commands, err := Parse(strings.NewReader(c))
assert.NoError(t, err)
commands2, err := Parse(strings.NewReader(Format(commands)))
assert.NoError(t, err)
assert.Equal(t, commands, commands2)
})
}
}