llama server wrapper
This commit is contained in:
Bruce MacDonald
parent
8fa91332fa
commit
0758cb2d4b
7 changed files with 83 additions and 138 deletions
34
server/README.md
Normal file
34
server/README.md
Normal file
|
|
@ -0,0 +1,34 @@
|
||||||
|
# Server
|
||||||
|
|
||||||
|
🙊
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
If using Apple silicon, you need a Python version that supports arm64:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-MacOSX-arm64.sh
|
||||||
|
bash Miniforge3-MacOSX-arm64.sh
|
||||||
|
```
|
||||||
|
|
||||||
|
Get the dependencies:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pip install llama-cpp-python
|
||||||
|
pip install -r requirements.txt
|
||||||
|
```
|
||||||
|
|
||||||
|
## Running
|
||||||
|
|
||||||
|
Put your model in `models/` and run:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python server.py
|
||||||
|
```
|
||||||
|
|
||||||
|
## API
|
||||||
|
|
||||||
|
### `POST /generate`
|
||||||
|
|
||||||
|
model: `string` - The name of the model to use in the `models` folder.
|
||||||
|
prompt: `string` - The prompt to use.
|
||||||
|
|
@ -1,2 +0,0 @@
|
||||||
LIBRARY_PATH=$PWD/go-llama.cpp C_INCLUDE_PATH=$PWD/go-llama.cpp go build .
|
|
||||||
|
|
||||||
|
|
@ -1,8 +0,0 @@
|
||||||
module github.com/keypairdev/keypair
|
|
||||||
|
|
||||||
go 1.20
|
|
||||||
|
|
||||||
require (
|
|
||||||
github.com/go-skynet/go-llama.cpp v0.0.0-20230620192753-7a36befaece1
|
|
||||||
github.com/sashabaranov/go-openai v1.11.3
|
|
||||||
)
|
|
||||||
|
|
@ -1,15 +0,0 @@
|
||||||
github.com/go-logr/logr v1.2.4 h1:g01GSCwiDw2xSZfjJ2/T9M+S6pFdcNtFYsp+Y43HYDQ=
|
|
||||||
github.com/go-skynet/go-llama.cpp v0.0.0-20230620192753-7a36befaece1 h1:UQ8y3kHxBgh3BnaW06y/X97fEN48yHPwWobMz8/aztU=
|
|
||||||
github.com/go-skynet/go-llama.cpp v0.0.0-20230620192753-7a36befaece1/go.mod h1:tzi97YvT1bVQ+iTG39LvpDkKG1WbizgtljC+orSoM40=
|
|
||||||
github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI=
|
|
||||||
github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
|
|
||||||
github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 h1:yAJXTCF9TqKcTiHJAE8dj7HMvPfh66eeA2JYW7eFpSE=
|
|
||||||
github.com/onsi/ginkgo/v2 v2.11.0 h1:WgqUCUt/lT6yXoQ8Wef0fsNn5cAuMK7+KT9UFRz2tcU=
|
|
||||||
github.com/onsi/gomega v1.27.8 h1:gegWiwZjBsf2DgiSbf5hpokZ98JVDMcWkUiigk6/KXc=
|
|
||||||
github.com/sashabaranov/go-openai v1.11.3 h1:bvwWF8hj4UhPlswBdL9/IfOpaHXfzGCJO8WY8ml9sGc=
|
|
||||||
github.com/sashabaranov/go-openai v1.11.3/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
|
|
||||||
golang.org/x/net v0.10.0 h1:X2//UzNDwYmtCLn7To6G58Wr6f5ahEAQgKNzv9Y951M=
|
|
||||||
golang.org/x/sys v0.9.0 h1:KS/R3tvhPqvJvwcKfnBHJwwthS11LRhmM5D59eEXa0s=
|
|
||||||
golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
|
|
||||||
golang.org/x/tools v0.9.3 h1:Gn1I8+64MsuTb/HpH+LmQtNas23LhUVr3rYZ0eKuaMM=
|
|
||||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
|
||||||
113
server/main.go
113
server/main.go
|
|
@ -1,113 +0,0 @@
|
||||||
package main
|
|
||||||
|
|
||||||
import (
|
|
||||||
"bytes"
|
|
||||||
"context"
|
|
||||||
"fmt"
|
|
||||||
"io"
|
|
||||||
"net/http"
|
|
||||||
"os"
|
|
||||||
"runtime"
|
|
||||||
|
|
||||||
"github.com/sashabaranov/go-openai"
|
|
||||||
|
|
||||||
llama "github.com/go-skynet/go-llama.cpp"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
type Model interface {
|
|
||||||
Name() string
|
|
||||||
Handler(w http.ResponseWriter, r *http.Request)
|
|
||||||
}
|
|
||||||
|
|
||||||
type LLama7B struct {
|
|
||||||
llama *llama.LLama
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewLLama7B() *LLama7B {
|
|
||||||
llama, err := llama.New("./models/7B/ggml-model-q4_0.bin", llama.EnableF16Memory, llama.SetContext(128), llama.EnableEmbeddings, llama.SetGPULayers(128))
|
|
||||||
if err != nil {
|
|
||||||
fmt.Println("Loading the model failed:", err.Error())
|
|
||||||
os.Exit(1)
|
|
||||||
}
|
|
||||||
|
|
||||||
return &LLama7B{
|
|
||||||
llama: llama,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (l *LLama7B) Name() string {
|
|
||||||
return "LLaMA 7B"
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *LLama7B) Handler(w http.ResponseWriter, r *http.Request) {
|
|
||||||
var text bytes.Buffer
|
|
||||||
io.Copy(&text, r.Body)
|
|
||||||
|
|
||||||
_, err := m.llama.Predict(text.String(), llama.Debug, llama.SetTokenCallback(func(token string) bool {
|
|
||||||
w.Write([]byte(token))
|
|
||||||
return true
|
|
||||||
}), llama.SetTokens(512), llama.SetThreads(runtime.NumCPU()), llama.SetTopK(90), llama.SetTopP(0.86), llama.SetStopWords("llama"))
|
|
||||||
|
|
||||||
if err != nil {
|
|
||||||
fmt.Println("Predict failed:", err.Error())
|
|
||||||
os.Exit(1)
|
|
||||||
}
|
|
||||||
|
|
||||||
embeds, err := m.llama.Embeddings(text.String())
|
|
||||||
if err != nil {
|
|
||||||
fmt.Printf("Embeddings: error %s \n", err.Error())
|
|
||||||
}
|
|
||||||
fmt.Printf("Embeddings: %v", embeds)
|
|
||||||
|
|
||||||
w.Header().Set("Content-Type", "text/event-stream")
|
|
||||||
w.Header().Set("Cache-Control", "no-cache")
|
|
||||||
w.Header().Set("Connection", "keep-alive")
|
|
||||||
}
|
|
||||||
|
|
||||||
type GPT4 struct {
|
|
||||||
apiKey string
|
|
||||||
}
|
|
||||||
|
|
||||||
func (g *GPT4) Name() string {
|
|
||||||
return "OpenAI GPT-4"
|
|
||||||
}
|
|
||||||
|
|
||||||
func (g *GPT4) Handler(w http.ResponseWriter, r *http.Request) {
|
|
||||||
w.WriteHeader(http.StatusOK)
|
|
||||||
client := openai.NewClient("your token")
|
|
||||||
resp, err := client.CreateChatCompletion(
|
|
||||||
context.Background(),
|
|
||||||
openai.ChatCompletionRequest{
|
|
||||||
Model: openai.GPT3Dot5Turbo,
|
|
||||||
Messages: []openai.ChatCompletionMessage{
|
|
||||||
{
|
|
||||||
Role: openai.ChatMessageRoleUser,
|
|
||||||
Content: "Hello!",
|
|
||||||
},
|
|
||||||
},
|
|
||||||
},
|
|
||||||
)
|
|
||||||
if err != nil {
|
|
||||||
fmt.Printf("chat completion error: %v\n", err)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
fmt.Println(resp.Choices[0].Message.Content)
|
|
||||||
|
|
||||||
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
|
||||||
w.WriteHeader(http.StatusOK)
|
|
||||||
}
|
|
||||||
|
|
||||||
// TODO: add subcommands to spawn different models
|
|
||||||
func main() {
|
|
||||||
model := &LLama7B{}
|
|
||||||
|
|
||||||
http.HandleFunc("/generate", model.Handler)
|
|
||||||
|
|
||||||
fmt.Println("Starting server on :8080")
|
|
||||||
if err := http.ListenAndServe(":8080", nil); err != nil {
|
|
||||||
fmt.Printf("Error starting server: %s\n", err)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
2
server/requirements.txt
Normal file
2
server/requirements.txt
Normal file
|
|
@ -0,0 +1,2 @@
|
||||||
|
Flask==2.3.2
|
||||||
|
flask_cors==3.0.10
|
||||||
47
server/server.py
Normal file
47
server/server.py
Normal file
|
|
@ -0,0 +1,47 @@
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
from llama_cpp import Llama
|
||||||
|
from flask import Flask, Response, stream_with_context, request
|
||||||
|
from flask_cors import CORS, cross_origin
|
||||||
|
|
||||||
|
app = Flask(__name__)
|
||||||
|
CORS(app) # enable CORS for all routes
|
||||||
|
|
||||||
|
# llms tracks which models are loaded
|
||||||
|
llms = {}
|
||||||
|
|
||||||
|
|
||||||
|
@app.route("/generate", methods=["POST"])
|
||||||
|
def generate():
|
||||||
|
data = request.get_json()
|
||||||
|
model = data.get("model")
|
||||||
|
prompt = data.get("prompt")
|
||||||
|
|
||||||
|
if not model:
|
||||||
|
return Response("Model is required", status=400)
|
||||||
|
if not prompt:
|
||||||
|
return Response("Prompt is required", status=400)
|
||||||
|
if not os.path.exists(f"../models/{model}.bin"):
|
||||||
|
return {"error": "The model file does not exist."}, 400
|
||||||
|
|
||||||
|
if model not in llms:
|
||||||
|
llms[model] = Llama(model_path=f"../models/{model}.bin")
|
||||||
|
|
||||||
|
def stream_response():
|
||||||
|
stream = llms[model](
|
||||||
|
str(prompt), # TODO: optimize prompt based on model
|
||||||
|
max_tokens=4096,
|
||||||
|
stop=["Q:", "\n"],
|
||||||
|
echo=True,
|
||||||
|
stream=True,
|
||||||
|
)
|
||||||
|
for output in stream:
|
||||||
|
yield json.dumps(output)
|
||||||
|
|
||||||
|
return Response(
|
||||||
|
stream_with_context(stream_response()), mimetype="text/event-stream"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
app.run(debug=True, threaded=True, port=5000)
|
||||||
Loading…
Reference in a new issue