mirror of
https://github.com/mudler/LocalAI.git
synced 2024-06-07 19:40:48 +00:00
354 lines
8.8 KiB
Go
354 lines
8.8 KiB
Go
package main
|
|
|
|
import (
|
|
"embed"
|
|
"fmt"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
|
|
model "github.com/go-skynet/llama-cli/pkg/model"
|
|
|
|
llama "github.com/go-skynet/go-llama.cpp"
|
|
"github.com/gofiber/fiber/v2"
|
|
"github.com/gofiber/fiber/v2/middleware/cors"
|
|
"github.com/gofiber/fiber/v2/middleware/filesystem"
|
|
"github.com/gofiber/fiber/v2/middleware/recover"
|
|
)
|
|
|
|
type OpenAIResponse struct {
|
|
Created int `json:"created,omitempty"`
|
|
Object string `json:"chat.completion,omitempty"`
|
|
ID string `json:"id,omitempty"`
|
|
Model string `json:"model,omitempty"`
|
|
Choices []Choice `json:"choices,omitempty"`
|
|
}
|
|
|
|
type Choice struct {
|
|
Index int `json:"index,omitempty"`
|
|
FinishReason string `json:"finish_reason,omitempty"`
|
|
Message Message `json:"message,omitempty"`
|
|
Text string `json:"text,omitempty"`
|
|
}
|
|
|
|
type Message struct {
|
|
Role string `json:"role,omitempty"`
|
|
Content string `json:"content,omitempty"`
|
|
}
|
|
|
|
type OpenAIModel struct {
|
|
ID string `json:"id"`
|
|
Object string `json:"object"`
|
|
}
|
|
|
|
//go:embed index.html
|
|
var indexHTML embed.FS
|
|
|
|
func completionEndpoint(defaultModel *llama.LLama, loader *model.ModelLoader, threads int, defaultMutex *sync.Mutex, mutexMap *sync.Mutex, mutexes map[string]*sync.Mutex) func(c *fiber.Ctx) error {
|
|
return func(c *fiber.Ctx) error {
|
|
|
|
var err error
|
|
var model *llama.LLama
|
|
|
|
// Get input data from the request body
|
|
input := new(struct {
|
|
Model string `json:"model"`
|
|
Prompt string `json:"prompt"`
|
|
})
|
|
if err := c.BodyParser(input); err != nil {
|
|
return err
|
|
}
|
|
|
|
if input.Model == "" {
|
|
if defaultModel == nil {
|
|
return fmt.Errorf("no default model loaded, and no model specified")
|
|
}
|
|
model = defaultModel
|
|
} else {
|
|
model, err = loader.LoadModel(input.Model)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// This is still needed, see: https://github.com/ggerganov/llama.cpp/discussions/784
|
|
if input.Model != "" {
|
|
mutexMap.Lock()
|
|
l, ok := mutexes[input.Model]
|
|
if !ok {
|
|
m := &sync.Mutex{}
|
|
mutexes[input.Model] = m
|
|
l = m
|
|
}
|
|
mutexMap.Unlock()
|
|
l.Lock()
|
|
defer l.Unlock()
|
|
} else {
|
|
defaultMutex.Lock()
|
|
defer defaultMutex.Unlock()
|
|
}
|
|
|
|
// Set the parameters for the language model prediction
|
|
topP, err := strconv.ParseFloat(c.Query("topP", "0.9"), 64) // Default value of topP is 0.9
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
topK, err := strconv.Atoi(c.Query("topK", "40")) // Default value of topK is 40
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
temperature, err := strconv.ParseFloat(c.Query("temperature", "0.5"), 64) // Default value of temperature is 0.5
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
tokens, err := strconv.Atoi(c.Query("tokens", "128")) // Default value of tokens is 128
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
predInput := input.Prompt
|
|
// A model can have a "file.bin.tmpl" file associated with a prompt template prefix
|
|
templatedInput, err := loader.TemplatePrefix(input.Model, struct {
|
|
Input string
|
|
}{Input: input.Prompt})
|
|
if err == nil {
|
|
predInput = templatedInput
|
|
}
|
|
|
|
// Generate the prediction using the language model
|
|
prediction, err := model.Predict(
|
|
predInput,
|
|
llama.SetTemperature(temperature),
|
|
llama.SetTopP(topP),
|
|
llama.SetTopK(topK),
|
|
llama.SetTokens(tokens),
|
|
llama.SetThreads(threads),
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Return the prediction in the response body
|
|
return c.JSON(OpenAIResponse{
|
|
Model: input.Model,
|
|
Choices: []Choice{{Text: prediction}},
|
|
})
|
|
}
|
|
}
|
|
|
|
func chatEndpoint(defaultModel *llama.LLama, loader *model.ModelLoader, threads int, defaultMutex *sync.Mutex, mutexMap *sync.Mutex, mutexes map[string]*sync.Mutex) func(c *fiber.Ctx) error {
|
|
return func(c *fiber.Ctx) error {
|
|
var err error
|
|
var model *llama.LLama
|
|
|
|
// Get input data from the request body
|
|
input := new(struct {
|
|
Messages []Message `json:"messages"`
|
|
Model string `json:"model"`
|
|
})
|
|
if err := c.BodyParser(input); err != nil {
|
|
return err
|
|
}
|
|
|
|
// TODO: drop me!
|
|
if input.Model == "gpt-3.5-turbo" {
|
|
input.Model = "ggml-koala-7b-model-q4_0-r2"
|
|
}
|
|
|
|
if input.Model == "" {
|
|
if defaultModel == nil {
|
|
return fmt.Errorf("no default model loaded, and no model specified")
|
|
}
|
|
model = defaultModel
|
|
} else {
|
|
model, err = loader.LoadModel(input.Model)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// This is still needed, see: https://github.com/ggerganov/llama.cpp/discussions/784
|
|
if input.Model != "" {
|
|
mutexMap.Lock()
|
|
l, ok := mutexes[input.Model]
|
|
if !ok {
|
|
m := &sync.Mutex{}
|
|
mutexes[input.Model] = m
|
|
l = m
|
|
}
|
|
mutexMap.Unlock()
|
|
l.Lock()
|
|
defer l.Unlock()
|
|
} else {
|
|
defaultMutex.Lock()
|
|
defer defaultMutex.Unlock()
|
|
}
|
|
|
|
// Set the parameters for the language model prediction
|
|
topP, err := strconv.ParseFloat(c.Query("topP", "0.9"), 64) // Default value of topP is 0.9
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
topK, err := strconv.Atoi(c.Query("topK", "40")) // Default value of topK is 40
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
temperature, err := strconv.ParseFloat(c.Query("temperature", "0.5"), 64) // Default value of temperature is 0.5
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
tokens, err := strconv.Atoi(c.Query("tokens", "128")) // Default value of tokens is 128
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
mess := []string{}
|
|
for _, i := range input.Messages {
|
|
mess = append(mess, i.Content)
|
|
}
|
|
|
|
predInput := strings.Join(mess, "\n")
|
|
|
|
// A model can have a "file.bin.tmpl" file associated with a prompt template prefix
|
|
templatedInput, err := loader.TemplatePrefix(input.Model, struct {
|
|
Input string
|
|
}{Input: predInput})
|
|
if err == nil {
|
|
predInput = templatedInput
|
|
}
|
|
|
|
// Generate the prediction using the language model
|
|
prediction, err := model.Predict(
|
|
predInput,
|
|
llama.SetTemperature(temperature),
|
|
llama.SetTopP(topP),
|
|
llama.SetTopK(topK),
|
|
llama.SetTokens(tokens),
|
|
llama.SetThreads(threads),
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Return the prediction in the response body
|
|
return c.JSON(OpenAIResponse{
|
|
Model: input.Model,
|
|
Choices: []Choice{{Message: Message{Role: "assistant", Content: prediction}}},
|
|
})
|
|
}
|
|
}
|
|
|
|
func Start(defaultModel *llama.LLama, loader *model.ModelLoader, listenAddr string, threads int) error {
|
|
app := fiber.New()
|
|
|
|
// Default middleware config
|
|
app.Use(recover.New())
|
|
app.Use(cors.New())
|
|
|
|
// This is still needed, see: https://github.com/ggerganov/llama.cpp/discussions/784
|
|
var mutex = &sync.Mutex{}
|
|
mu := map[string]*sync.Mutex{}
|
|
var mumutex = &sync.Mutex{}
|
|
|
|
// openAI compatible API endpoint
|
|
app.Post("/v1/chat/completions", chatEndpoint(defaultModel, loader, threads, mutex, mumutex, mu))
|
|
app.Post("/v1/completions", completionEndpoint(defaultModel, loader, threads, mutex, mumutex, mu))
|
|
app.Get("/v1/models", func(c *fiber.Ctx) error {
|
|
models, err := loader.ListModels()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
dataModels := []OpenAIModel{}
|
|
for _, m := range models {
|
|
dataModels = append(dataModels, OpenAIModel{ID: m, Object: "model"})
|
|
}
|
|
return c.JSON(struct {
|
|
Object string `json:"object"`
|
|
Data []OpenAIModel `json:"data"`
|
|
}{
|
|
Object: "list",
|
|
Data: dataModels,
|
|
})
|
|
})
|
|
|
|
app.Use("/", filesystem.New(filesystem.Config{
|
|
Root: http.FS(indexHTML),
|
|
NotFoundFile: "index.html",
|
|
}))
|
|
|
|
/*
|
|
curl --location --request POST 'http://localhost:8080/predict' --header 'Content-Type: application/json' --data-raw '{
|
|
"text": "What is an alpaca?",
|
|
"topP": 0.8,
|
|
"topK": 50,
|
|
"temperature": 0.7,
|
|
"tokens": 100
|
|
}'
|
|
*/
|
|
// Endpoint to generate the prediction
|
|
app.Post("/predict", func(c *fiber.Ctx) error {
|
|
mutex.Lock()
|
|
defer mutex.Unlock()
|
|
// Get input data from the request body
|
|
input := new(struct {
|
|
Text string `json:"text"`
|
|
})
|
|
if err := c.BodyParser(input); err != nil {
|
|
return err
|
|
}
|
|
|
|
// Set the parameters for the language model prediction
|
|
topP, err := strconv.ParseFloat(c.Query("topP", "0.9"), 64) // Default value of topP is 0.9
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
topK, err := strconv.Atoi(c.Query("topK", "40")) // Default value of topK is 40
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
temperature, err := strconv.ParseFloat(c.Query("temperature", "0.5"), 64) // Default value of temperature is 0.5
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
tokens, err := strconv.Atoi(c.Query("tokens", "128")) // Default value of tokens is 128
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Generate the prediction using the language model
|
|
prediction, err := defaultModel.Predict(
|
|
input.Text,
|
|
llama.SetTemperature(temperature),
|
|
llama.SetTopP(topP),
|
|
llama.SetTopK(topK),
|
|
llama.SetTokens(tokens),
|
|
llama.SetThreads(threads),
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Return the prediction in the response body
|
|
return c.JSON(struct {
|
|
Prediction string `json:"prediction"`
|
|
}{
|
|
Prediction: prediction,
|
|
})
|
|
})
|
|
|
|
// Start the server
|
|
app.Listen(listenAddr)
|
|
return nil
|
|
}
|