routes.go 3.34 KB
Newer Older
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1
2
3
package server

import (
Michael Yang's avatar
Michael Yang committed
4
	"embed"
Michael Yang's avatar
Michael Yang committed
5
	"encoding/json"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
6
7
8
	"fmt"
	"io"
	"log"
Michael Yang's avatar
Michael Yang committed
9
	"math"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
10
11
	"net"
	"net/http"
Michael Yang's avatar
Michael Yang committed
12
	"path"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
13
	"runtime"
Michael Yang's avatar
Michael Yang committed
14
15
	"strings"
	"text/template"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
16
17

	"github.com/gin-gonic/gin"
Michael Yang's avatar
Michael Yang committed
18
	"github.com/lithammer/fuzzysearch/fuzzy"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
19

Jeffrey Morgan's avatar
Jeffrey Morgan committed
20
	"github.com/jmorganca/ollama/api"
Michael Yang's avatar
Michael Yang committed
21
	"github.com/jmorganca/ollama/llama"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
22
23
)

Michael Yang's avatar
Michael Yang committed
24
25
26
//go:embed templates/*
var templatesFS embed.FS
var templates = template.Must(template.ParseFS(templatesFS, "templates/*.prompt"))
Michael Yang's avatar
Michael Yang committed
27

Bruce MacDonald's avatar
Bruce MacDonald committed
28
func generate(c *gin.Context) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
29
	// TODO: these should be request parameters
Michael Yang's avatar
Michael Yang committed
30
	gpulayers := 1
Jeffrey Morgan's avatar
Jeffrey Morgan committed
31
32
33
	tokens := 512
	threads := runtime.NumCPU()

Bruce MacDonald's avatar
Bruce MacDonald committed
34
35
36
37
38
	var req api.GenerateRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
		return
	}
39

Michael Yang's avatar
Michael Yang committed
40
41
42
43
	if remoteModel, _ := getRemote(req.Model); remoteModel != nil {
		req.Model = remoteModel.FullName()
	}

Michael Yang's avatar
Michael Yang committed
44
	model, err := llama.New(req.Model, llama.EnableF16Memory, llama.SetContext(128), llama.EnableEmbeddings, llama.SetGPULayers(gpulayers))
Bruce MacDonald's avatar
Bruce MacDonald committed
45
46
47
48
	if err != nil {
		fmt.Println("Loading the model failed:", err.Error())
		return
	}
Michael Yang's avatar
Michael Yang committed
49
	defer model.Free()
Jeffrey Morgan's avatar
Jeffrey Morgan committed
50

Michael Yang's avatar
Michael Yang committed
51
52
53
54
55
	templateNames := make([]string, 0, len(templates.Templates()))
	for _, template := range templates.Templates() {
		templateNames = append(templateNames, template.Name())
	}

Michael Yang's avatar
Michael Yang committed
56
	match, _ := matchRankOne(path.Base(req.Model), templateNames)
Michael Yang's avatar
Michael Yang committed
57
58
59
60
61
62
63
64
65
66
	if template := templates.Lookup(match); template != nil {
		var sb strings.Builder
		if err := template.Execute(&sb, req); err != nil {
			fmt.Println("Prompt template failed:", err.Error())
			return
		}

		req.Prompt = sb.String()
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
67
	ch := make(chan string)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
68

Bruce MacDonald's avatar
Bruce MacDonald committed
69
70
	go func() {
		defer close(ch)
Michael Yang's avatar
Michael Yang committed
71
		_, err := model.Predict(req.Prompt, llama.Debug, llama.SetTokenCallback(func(token string) bool {
Bruce MacDonald's avatar
Bruce MacDonald committed
72
73
74
			ch <- token
			return true
		}), llama.SetTokens(tokens), llama.SetThreads(threads), llama.SetTopK(90), llama.SetTopP(0.86), llama.SetStopWords("llama"))
Jeffrey Morgan's avatar
Jeffrey Morgan committed
75
		if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
76
			panic(err)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
77
		}
Bruce MacDonald's avatar
Bruce MacDonald committed
78
	}()
Jeffrey Morgan's avatar
Jeffrey Morgan committed
79

Bruce MacDonald's avatar
Bruce MacDonald committed
80
	c.Stream(func(w io.Writer) bool {
Michael Yang's avatar
Michael Yang committed
81
		token, ok := <-ch
Bruce MacDonald's avatar
Bruce MacDonald committed
82
83
84
		if !ok {
			return false
		}
Michael Yang's avatar
Michael Yang committed
85

Michael Yang's avatar
Michael Yang committed
86
87
		resp := api.GenerateResponse{
			Response: token,
Michael Yang's avatar
Michael Yang committed
88
89
90
91
92
93
94
95
96
97
98
99
		}

		bts, err := json.Marshal(resp)
		if err != nil {
			return false
		}

		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
			return false
		}

Bruce MacDonald's avatar
Bruce MacDonald committed
100
		return true
Jeffrey Morgan's avatar
Jeffrey Morgan committed
101
	})
Bruce MacDonald's avatar
Bruce MacDonald committed
102
103
104
105
106
}

func Serve(ln net.Listener) error {
	r := gin.Default()

Bruce MacDonald's avatar
Bruce MacDonald committed
107
108
109
110
111
112
113
	r.POST("api/pull", func(c *gin.Context) {
		var req api.PullRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
			return
		}

Bruce MacDonald's avatar
Bruce MacDonald committed
114
		progressCh := make(chan api.PullProgress)
Bruce MacDonald's avatar
Bruce MacDonald committed
115
116
117
118
119
120
121
122
123
124
125
126
127
		go func() {
			defer close(progressCh)
			if err := pull(req.Model, progressCh); err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
				return
			}
		}()

		c.Stream(func(w io.Writer) bool {
			progress, ok := <-progressCh
			if !ok {
				return false
			}
Michael Yang's avatar
Michael Yang committed
128
129
130
131
132
133
134
135
136
137
138

			bts, err := json.Marshal(progress)
			if err != nil {
				return false
			}

			bts = append(bts, '\n')
			if _, err := w.Write(bts); err != nil {
				return false
			}

Bruce MacDonald's avatar
Bruce MacDonald committed
139
140
141
			return true
		})
	})
Bruce MacDonald's avatar
Bruce MacDonald committed
142
143

	r.POST("/api/generate", generate)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
144
145
146
147
148
149
150
151

	log.Printf("Listening on %s", ln.Addr())
	s := &http.Server{
		Handler: r,
	}

	return s.Serve(ln)
}
Michael Yang's avatar
Michael Yang committed
152
153

func matchRankOne(source string, targets []string) (bestMatch string, bestRank int) {
Michael Yang's avatar
Michael Yang committed
154
	bestRank = math.MaxInt
Michael Yang's avatar
Michael Yang committed
155
	for _, target := range targets {
Michael Yang's avatar
Michael Yang committed
156
		if rank := fuzzy.LevenshteinDistance(source, target); bestRank > rank {
Michael Yang's avatar
Michael Yang committed
157
158
159
160
161
162
163
			bestRank = rank
			bestMatch = target
		}
	}

	return
}