mirror of https://gitee.com/namelin2022/ollama
Devon Rifkin
1 year ago
committed by
GitHub
GitHub
6 changed files with 281 additions and 269 deletions
@ -0,0 +1,134 @@ |
|||
package thinking |
|||
|
|||
import ( |
|||
"strings" |
|||
"text/template" |
|||
"text/template/parse" |
|||
) |
|||
|
|||
func templateVisit(n parse.Node, enterFn func(parse.Node) bool, exitFn func(parse.Node)) { |
|||
if n == nil { |
|||
return |
|||
} |
|||
shouldContinue := enterFn(n) |
|||
if !shouldContinue { |
|||
return |
|||
} |
|||
switch x := n.(type) { |
|||
case *parse.ListNode: |
|||
for _, c := range x.Nodes { |
|||
templateVisit(c, enterFn, exitFn) |
|||
} |
|||
case *parse.BranchNode: |
|||
if x.Pipe != nil { |
|||
templateVisit(x.Pipe, enterFn, exitFn) |
|||
} |
|||
if x.List != nil { |
|||
templateVisit(x.List, enterFn, exitFn) |
|||
} |
|||
if x.ElseList != nil { |
|||
templateVisit(x.ElseList, enterFn, exitFn) |
|||
} |
|||
case *parse.ActionNode: |
|||
templateVisit(x.Pipe, enterFn, exitFn) |
|||
case *parse.WithNode: |
|||
templateVisit(&x.BranchNode, enterFn, exitFn) |
|||
case *parse.RangeNode: |
|||
templateVisit(&x.BranchNode, enterFn, exitFn) |
|||
case *parse.IfNode: |
|||
templateVisit(&x.BranchNode, enterFn, exitFn) |
|||
case *parse.TemplateNode: |
|||
templateVisit(x.Pipe, enterFn, exitFn) |
|||
case *parse.PipeNode: |
|||
for _, c := range x.Cmds { |
|||
templateVisit(c, enterFn, exitFn) |
|||
} |
|||
case *parse.CommandNode: |
|||
for _, a := range x.Args { |
|||
templateVisit(a, enterFn, exitFn) |
|||
} |
|||
// text, field, number, etc. are leaves – nothing to recurse into
|
|||
} |
|||
if exitFn != nil { |
|||
exitFn(n) |
|||
} |
|||
} |
|||
|
|||
// InferTags uses a heuristic to infer the tags that surround thinking traces:
|
|||
// We look for a range node that iterates over "Messages" and then look for a
|
|||
// reference to "Thinking" like `{{.Thinking}}`. We then go up to the nearest
|
|||
// ListNode and take the first and last TextNodes as the opening and closing
|
|||
// tags.
|
|||
func InferTags(t *template.Template) (string, string) { |
|||
ancestors := []parse.Node{} |
|||
|
|||
openingTag := "" |
|||
closingTag := "" |
|||
|
|||
enterFn := func(n parse.Node) bool { |
|||
ancestors = append(ancestors, n) |
|||
|
|||
switch x := n.(type) { |
|||
case *parse.FieldNode: |
|||
if len(x.Ident) > 0 && x.Ident[0] == "Thinking" { |
|||
var mostRecentRange *parse.RangeNode |
|||
for i := len(ancestors) - 1; i >= 0; i-- { |
|||
if r, ok := ancestors[i].(*parse.RangeNode); ok { |
|||
mostRecentRange = r |
|||
break |
|||
} |
|||
} |
|||
if mostRecentRange == nil || !rangeUsesField(mostRecentRange, "Messages") { |
|||
return true |
|||
} |
|||
|
|||
// TODO(drifkin): to be more robust, check that it's in the action
|
|||
// part, not the `if`'s pipeline part. We do match on the nearest list
|
|||
// that starts and ends with text nodes, which makes this not strictly
|
|||
// necessary for our heuristic
|
|||
|
|||
// go up to the nearest ancestor that is a *parse.ListNode
|
|||
for i := len(ancestors) - 1; i >= 0; i-- { |
|||
if l, ok := ancestors[i].(*parse.ListNode); ok { |
|||
firstNode := l.Nodes[0] |
|||
if t, ok := firstNode.(*parse.TextNode); ok { |
|||
openingTag = strings.TrimSpace(t.String()) |
|||
} |
|||
lastNode := l.Nodes[len(l.Nodes)-1] |
|||
if t, ok := lastNode.(*parse.TextNode); ok { |
|||
closingTag = strings.TrimSpace(t.String()) |
|||
} |
|||
|
|||
break |
|||
} |
|||
} |
|||
} |
|||
} |
|||
|
|||
return true |
|||
} |
|||
|
|||
exitFn := func(n parse.Node) { |
|||
ancestors = ancestors[:len(ancestors)-1] |
|||
} |
|||
|
|||
templateVisit(t.Root, enterFn, exitFn) |
|||
|
|||
return openingTag, closingTag |
|||
} |
|||
|
|||
// checks to see if the given field name is present in the pipeline of the given range node
|
|||
func rangeUsesField(rangeNode *parse.RangeNode, field string) bool { |
|||
found := false |
|||
enterFn := func(n parse.Node) bool { |
|||
switch x := n.(type) { |
|||
case *parse.FieldNode: |
|||
if x.Ident[0] == field { |
|||
found = true |
|||
} |
|||
} |
|||
return true |
|||
} |
|||
templateVisit(rangeNode.BranchNode.Pipe, enterFn, nil) |
|||
return found |
|||
} |
|||
@ -0,0 +1,130 @@ |
|||
package thinking |
|||
|
|||
import ( |
|||
"testing" |
|||
"text/template" |
|||
) |
|||
|
|||
func TestInferThinkingTags(t *testing.T) { |
|||
cases := []struct { |
|||
desc string |
|||
tmplString string |
|||
wantOpeningTag string |
|||
wantClosingTag string |
|||
}{ |
|||
{ |
|||
desc: "basic", |
|||
tmplString: ` |
|||
{{ if .Thinking}} |
|||
/think |
|||
{{ end }} |
|||
{{- range $i, $_ := .Messages }} |
|||
{{- $last := eq (len (slice $.Messages $i)) 1 -}} |
|||
{{ if and $last .Thinking }} |
|||
<think>{{ .Thinking }}</think> |
|||
{{ end }} |
|||
{{ end }} |
|||
`, |
|||
wantOpeningTag: "<think>", |
|||
wantClosingTag: "</think>", |
|||
}, |
|||
{ |
|||
desc: "doubly nested range", |
|||
tmplString: ` |
|||
{{ if .Thinking}} |
|||
/think |
|||
{{ end }} |
|||
{{- range $i, $_ := .Messages }} |
|||
{{- range $j, $_ := .NotMessages }} |
|||
{{- $last := eq (len (slice $.Messages $i)) 1 -}} |
|||
{{ if and $last .Thinking }} |
|||
<think>{{ .Thinking }}</think> |
|||
{{ end }} |
|||
{{ end }} |
|||
{{ end }} |
|||
`, |
|||
wantOpeningTag: "", |
|||
wantClosingTag: "", |
|||
}, |
|||
{ |
|||
desc: "whitespace is trimmed", |
|||
tmplString: ` |
|||
{{ if .Thinking}} |
|||
/think |
|||
{{ end }} |
|||
{{- range $i, $_ := .Messages }} |
|||
{{- $last := eq (len (slice $.Messages $i)) 1 -}} |
|||
{{ if and $last .Thinking }} |
|||
Some text before {{ .Thinking }} Some text after |
|||
{{ end }} |
|||
{{ end }} |
|||
`, |
|||
wantOpeningTag: "Some text before", |
|||
wantClosingTag: "Some text after", |
|||
}, |
|||
{ |
|||
desc: "qwen3", |
|||
tmplString: ` |
|||
{{- if or .System .Tools .Thinking }}<|im_start|>system |
|||
{{- if .System }} |
|||
{{ .System }} |
|||
{{- end }} |
|||
{{- if .Tools }} |
|||
|
|||
# Tools |
|||
|
|||
You may call one or more functions to assist with the user query. |
|||
|
|||
You are provided with function signatures within <tools></tools> XML tags: |
|||
<tools> |
|||
{{- range .Tools }} |
|||
{"type": "function", "function": {{ .Function }}} |
|||
{{- end }} |
|||
</tools> |
|||
|
|||
For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags: |
|||
<tool_call> |
|||
{"name": <function-name>, "arguments": <args-json-object>} |
|||
</tool_call> |
|||
{{- end }} |
|||
{{- if .Thinking }} |
|||
/think |
|||
{{- else }} |
|||
/no_think |
|||
{{- end }}<|im_end|> |
|||
{{ end }} |
|||
{{- range $i, $_ := .Messages }} |
|||
{{- $last := eq (len (slice $.Messages $i)) 1 -}} |
|||
{{- if eq .Role "user" }}<|im_start|>user |
|||
{{ .Content }}<|im_end|> |
|||
{{ else if eq .Role "assistant" }}<|im_start|>assistant |
|||
{{ if and $last .Thinking }} |
|||
<think>{{ .Thinking }}</think> |
|||
{{ end }} |
|||
{{ if .Content }}{{ .Content }} |
|||
{{- else if .ToolCalls }}<tool_call> |
|||
{{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}} |
|||
{{ end }}</tool_call> |
|||
{{- end }}{{ if not $last }}<|im_end|> |
|||
{{ end }} |
|||
{{- else if eq .Role "tool" }}<|im_start|>user |
|||
<tool_response> |
|||
{{ .Content }} |
|||
</tool_response><|im_end|> |
|||
{{ end }} |
|||
{{- if and (ne .Role "assistant") $last }}<|im_start|>assistant |
|||
{{ end }} |
|||
{{- end }} |
|||
`, |
|||
wantOpeningTag: "<think>", |
|||
wantClosingTag: "</think>", |
|||
}, |
|||
} |
|||
for _, c := range cases { |
|||
tmpl := template.Must(template.New("test").Parse(c.tmplString)) |
|||
openingTag, closingTag := InferTags(tmpl) |
|||
if openingTag != c.wantOpeningTag || closingTag != c.wantClosingTag { |
|||
t.Errorf("case %q: got (%q,%q), want (%q,%q)", c.desc, openingTag, closingTag, c.wantOpeningTag, c.wantClosingTag) |
|||
} |
|||
} |
|||
} |
|||
Loading…
Reference in new issue