third-party-mirrors/ollama
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

enable --reranking flag for rerank handler while starting server

Browse Source 
Signed-off-by: Liu Yuan <namei.unix@gmail.com>
This commit is contained in:
Liu Yuan 2024-10-31 22:48:27 +08:00
parent 67818b5093
commit bfae776f34
3 changed files with 9 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
api/types.go
View File

@ -242,6 +242,7 @@ type Runner struct {
UseMMap *bool `json:"use_mmap,omitempty"` UseMMap *bool `json:"use_mmap,omitempty"`
UseMLock bool `json:"use_mlock,omitempty"` UseMLock bool `json:"use_mlock,omitempty"`
NumThread int `json:"num_thread,omitempty"` NumThread int `json:"num_thread,omitempty"`
Reranking bool `json:"reranking,omitempty"`
} }
// EmbedRequest is the request passed to [Client.Embed]. // EmbedRequest is the request passed to [Client.Embed].

4
llm/server.go
View File

@ -189,6 +189,10 @@ func NewLlamaServer(gpus discover.GpuInfoList, model string, ggml *GGML, adapter
"--batch-size", strconv.Itoa(opts.NumBatch), "--batch-size", strconv.Itoa(opts.NumBatch),
} }
if opts.Reranking {
params = append(params, "--reranking")
}
if opts.NumGPU >= 0 { if opts.NumGPU >= 0 {
params = append(params, "--n-gpu-layers", strconv.Itoa(opts.NumGPU)) params = append(params, "--n-gpu-layers", strconv.Itoa(opts.NumGPU))
} }

4
server/routes.go
View File

@ -371,6 +371,10 @@ func (s *Server) RerankHandler(c *gin.Context) {
return return
} }
if req.Options == nil {
req.Options = make(map[string]any)
}
req.Options["reranking"] = true
r, _, _, err := s.scheduleRunner(c.Request.Context(), req.Model, []Capability{}, req.Options, req.KeepAlive) r, _, _, err := s.scheduleRunner(c.Request.Context(), req.Model, []Capability{}, req.Options, req.KeepAlive)
if err != nil { if err != nil {
handleScheduleError(c, req.Model, err) handleScheduleError(c, req.Model, err)