blob: 4690563578e3099a7371f5be3dcc69725e80266c (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
|
package llm
import (
"fmt"
"strings"
)
type Session struct {
Messages []ChatMLer `json:"messages"`
}
type ChatMLer interface {
ChatML() string
}
func (s Session) ChatML() string {
var sb strings.Builder
for _, message := range s.Messages {
fmt.Fprintf(&sb, "%s\n", message.ChatML())
}
return sb.String()
}
type Message struct {
Role string `json:"role"`
Content string `json:"content"`
}
func (m Message) ChatML() string {
if m.Content == "" {
return fmt.Sprintf("<|im_start|>%s\n", m.Role)
}
return fmt.Sprintf("<|im_start|>%s\n%s\n<|im_end|>", m.Role, m.Content)
}
|
