ollama/x/model/file.go
Blake Mizerany 7c7f56a7fb x/model: limit part len, not entire len
Limiting the whole name length comes naturally with part name length
restrictions. This aligns with Docker's registry behavior.
2024-04-05 18:49:11 -07:00

133 lines
2.5 KiB
Go

// Package model implements the File and Name types for working with and
// representing Modelfiles and model Names.
//
// The Name type should be used when working with model names, and the File
// type should be used when working with Modelfiles.
package model
import (
"bufio"
"io"
"iter"
"strings"
)
type ParamPragma struct {
Key string
Value string
}
type MessagePragma struct {
Role string
Content string
}
type File struct {
// From is a required pragma that specifies the source of the model,
// either on disk, or by reference (see model.ParseName).
From string
// Optional
Params []ParamPragma
Template string
System string
Adapter string
Messages []MessagePragma
License string
}
type FileError struct {
Pragma string
Message string
}
func (e *FileError) Error() string {
return e.Pragma + ": " + e.Message
}
// Pragma represents a single pragma in a Modelfile.
type Pragma struct {
// The pragma name
Name string
// Args contains the user-defined arguments for the pragma. If no
// arguments were provided, it is nil.
Args []string
}
func (p Pragma) Arg(i int) string {
if i >= len(p.Args) {
return ""
}
return p.Args[i]
}
func FilePragmas(r io.Reader) iter.Seq2[Pragma, error] {
return func(yield func(Pragma, error) bool) {
sc := bufio.NewScanner(r)
for sc.Scan() {
line := sc.Text()
// TODO(bmizerany): set a max num fields/args to
// prevent mem bloat
args := strings.Fields(line)
if len(args) == 0 {
continue
}
p := Pragma{
Name: strings.ToUpper(args[0]),
}
if p.Name == "MESSAGE" {
// handle special case where message content
// is space separated on the _rest_ of the
// line like: `MESSAGE user Is Ontario in
// Canada?`
panic("TODO")
}
if len(args) > 1 {
p.Args = args[1:]
}
if !yield(p, nil) {
return
}
}
if sc.Err() != nil {
yield(Pragma{}, sc.Err())
}
}
}
func ParseFile(r io.Reader) (File, error) {
var f File
for p, err := range FilePragmas(r) {
if err != nil {
return File{}, err
}
switch p.Name {
case "FROM":
f.From = p.Arg(0)
case "PARAMETER":
f.Params = append(f.Params, ParamPragma{
Key: strings.ToLower(p.Arg(0)),
Value: p.Arg(1),
})
case "TEMPLATE":
f.Template = p.Arg(0)
case "SYSTEM":
f.System = p.Arg(0)
case "ADAPTER":
f.Adapter = p.Arg(0)
case "MESSAGE":
f.Messages = append(f.Messages, MessagePragma{
Role: p.Arg(0),
Content: p.Arg(1),
})
case "LICENSE":
f.License = p.Arg(0)
}
}
return f, nil
}