diff --git a/parser/parser.go b/parser/parser.go index 686a1e69..7f566da4 100644 --- a/parser/parser.go +++ b/parser/parser.go @@ -124,7 +124,7 @@ func ParseFile(r io.Reader) (*File, error) { case stateComment, stateNil: // pass case stateValue: - s, ok := unquote(b.String()) + s, ok := unquote(strings.TrimSpace(b.String())) if !ok || isSpace(r) { if _, err := b.WriteRune(r); err != nil { return nil, err @@ -158,7 +158,7 @@ func ParseFile(r io.Reader) (*File, error) { case stateComment, stateNil: // pass; nothing to flush case stateValue: - s, ok := unquote(b.String()) + s, ok := unquote(strings.TrimSpace(b.String())) if !ok { return nil, io.ErrUnexpectedEOF } diff --git a/parser/parser_test.go b/parser/parser_test.go index 7123e53b..2b5c4c88 100644 --- a/parser/parser_test.go +++ b/parser/parser_test.go @@ -22,7 +22,13 @@ ADAPTER adapter1 LICENSE MIT PARAMETER param1 value1 PARAMETER param2 value2 -TEMPLATE template1 +TEMPLATE """{{ if .System }}<|start_header_id|>system<|end_header_id|> + +{{ .System }}<|eot_id|>{{ end }}{{ if .Prompt }}<|start_header_id|>user<|end_header_id|> + +{{ .Prompt }}<|eot_id|>{{ end }}<|start_header_id|>assistant<|end_header_id|> + +{{ .Response }}<|eot_id|>""" ` reader := strings.NewReader(input) @@ -36,7 +42,40 @@ TEMPLATE template1 {Name: "license", Args: "MIT"}, {Name: "param1", Args: "value1"}, {Name: "param2", Args: "value2"}, - {Name: "template", Args: "template1"}, + {Name: "template", Args: "{{ if .System }}<|start_header_id|>system<|end_header_id|>\n\n{{ .System }}<|eot_id|>{{ end }}{{ if .Prompt }}<|start_header_id|>user<|end_header_id|>\n\n{{ .Prompt }}<|eot_id|>{{ end }}<|start_header_id|>assistant<|end_header_id|>\n\n{{ .Response }}<|eot_id|>"}, + } + + assert.Equal(t, expectedCommands, modelfile.Commands) +} + +func TestParseFileTrimSpace(t *testing.T) { + input := ` +FROM " model 1" +ADAPTER adapter3 +LICENSE "MIT " +PARAMETER param1 value1 +PARAMETER param2 value2 +TEMPLATE """ {{ if .System }}<|start_header_id|>system<|end_header_id|> + +{{ .System }}<|eot_id|>{{ end }}{{ if .Prompt }}<|start_header_id|>user<|end_header_id|> + +{{ .Prompt }}<|eot_id|>{{ end }}<|start_header_id|>assistant<|end_header_id|> + +{{ .Response }}<|eot_id|> """ +` + + reader := strings.NewReader(input) + + modelfile, err := ParseFile(reader) + require.NoError(t, err) + + expectedCommands := []Command{ + {Name: "model", Args: " model 1"}, + {Name: "adapter", Args: "adapter3"}, + {Name: "license", Args: "MIT "}, + {Name: "param1", Args: "value1"}, + {Name: "param2", Args: "value2"}, + {Name: "template", Args: " {{ if .System }}<|start_header_id|>system<|end_header_id|>\n\n{{ .System }}<|eot_id|>{{ end }}{{ if .Prompt }}<|start_header_id|>user<|end_header_id|>\n\n{{ .Prompt }}<|eot_id|>{{ end }}<|start_header_id|>assistant<|end_header_id|>\n\n{{ .Response }}<|eot_id|> "}, } assert.Equal(t, expectedCommands, modelfile.Commands) @@ -48,6 +87,26 @@ func TestParseFileFrom(t *testing.T) { expected []Command err error }{ + { + "FROM \"FOO BAR \"", + []Command{{Name: "model", Args: "FOO BAR "}}, + nil, + }, + { + "FROM \"FOO BAR\"\nPARAMETER param1 value1", + []Command{{Name: "model", Args: "FOO BAR"}, {Name: "param1", Args: "value1"}}, + nil, + }, + { + "FROM FOOO BAR ", + []Command{{Name: "model", Args: "FOOO BAR"}}, + nil, + }, + { + "FROM /what/is/the path ", + []Command{{Name: "model", Args: "/what/is/the path"}}, + nil, + }, { "FROM foo", []Command{{Name: "model", Args: "foo"}}, @@ -86,6 +145,11 @@ func TestParseFileFrom(t *testing.T) { []Command{{Name: "param1", Args: "value1"}, {Name: "model", Args: "foo"}}, nil, }, + { + "PARAMETER what the \nFROM lemons make lemonade ", + []Command{{Name: "what", Args: "the"}, {Name: "model", Args: "lemons make lemonade"}}, + nil, + }, } for _, c := range cases { @@ -399,7 +463,7 @@ func TestParseFileParameters(t *testing.T) { "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:\"\"\"": {"stop", "### User:"},