run prompts
This commit is contained in:
parent
7d500692f2
commit
3d6009aae3
4 changed files with 91 additions and 32 deletions
|
@ -5,6 +5,7 @@ import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
@ -63,20 +64,19 @@ func (c *Client) stream(ctx context.Context, method string, path string, reqData
|
||||||
|
|
||||||
for {
|
for {
|
||||||
line, err := reader.ReadBytes('\n')
|
line, err := reader.ReadBytes('\n')
|
||||||
if err != nil {
|
switch {
|
||||||
if err == io.EOF {
|
case errors.Is(err, io.EOF):
|
||||||
break
|
return nil
|
||||||
} else {
|
case err != nil:
|
||||||
return err // Handle other errors
|
return err
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := checkError(res, line); err != nil {
|
if err := checkError(res, line); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
callback(bytes.TrimSuffix(line, []byte("\n")))
|
callback(bytes.TrimSuffix(line, []byte("\n")))
|
||||||
}
|
}
|
||||||
|
|
||||||
return nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *Client) do(ctx context.Context, method string, path string, reqData any, respData any) error {
|
func (c *Client) do(ctx context.Context, method string, path string, reqData any, respData any) error {
|
||||||
|
@ -124,11 +124,9 @@ func (c *Client) do(ctx context.Context, method string, path string, reqData any
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *Client) Generate(ctx context.Context, req *GenerateRequest, callback func(token string)) (*GenerateResponse, error) {
|
func (c *Client) Generate(ctx context.Context, req *GenerateRequest, callback func(bts []byte)) (*GenerateResponse, error) {
|
||||||
var res GenerateResponse
|
var res GenerateResponse
|
||||||
if err := c.stream(ctx, http.MethodPost, "/api/generate", req, func(token []byte) {
|
if err := c.stream(ctx, http.MethodPost, "/api/generate", req, callback); err != nil {
|
||||||
callback(string(token))
|
|
||||||
}); err != nil {
|
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
98
cmd/cmd.go
98
cmd/cmd.go
|
@ -1,7 +1,9 @@
|
||||||
package cmd
|
package cmd
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bufio"
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"log"
|
"log"
|
||||||
"net"
|
"net"
|
||||||
|
@ -10,9 +12,11 @@ import (
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
"github.com/gosuri/uiprogress"
|
"github.com/gosuri/uiprogress"
|
||||||
|
"github.com/spf13/cobra"
|
||||||
|
"golang.org/x/term"
|
||||||
|
|
||||||
"github.com/jmorganca/ollama/api"
|
"github.com/jmorganca/ollama/api"
|
||||||
"github.com/jmorganca/ollama/server"
|
"github.com/jmorganca/ollama/server"
|
||||||
"github.com/spf13/cobra"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
func cacheDir() string {
|
func cacheDir() string {
|
||||||
|
@ -28,13 +32,13 @@ func bytesToGB(bytes int) float64 {
|
||||||
return float64(bytes) / float64(1<<30)
|
return float64(bytes) / float64(1<<30)
|
||||||
}
|
}
|
||||||
|
|
||||||
func run(model string) error {
|
func RunRun(cmd *cobra.Command, args []string) error {
|
||||||
client, err := NewAPIClient()
|
client, err := NewAPIClient()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
pr := api.PullRequest{
|
pr := api.PullRequest{
|
||||||
Model: model,
|
Model: args[0],
|
||||||
}
|
}
|
||||||
var bar *uiprogress.Bar
|
var bar *uiprogress.Bar
|
||||||
mutex := &sync.Mutex{}
|
mutex := &sync.Mutex{}
|
||||||
|
@ -60,10 +64,71 @@ func run(model string) error {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
fmt.Println("Up to date.")
|
fmt.Println("Up to date.")
|
||||||
|
return RunGenerate(cmd, args)
|
||||||
|
}
|
||||||
|
|
||||||
|
func RunGenerate(_ *cobra.Command, args []string) error {
|
||||||
|
if len(args) > 1 {
|
||||||
|
return generate(args[0], args[1:]...)
|
||||||
|
}
|
||||||
|
|
||||||
|
if term.IsTerminal(int(os.Stdin.Fd())) {
|
||||||
|
return generateInteractive(args[0])
|
||||||
|
}
|
||||||
|
|
||||||
|
return generateBatch(args[0])
|
||||||
|
}
|
||||||
|
|
||||||
|
func generate(model string, prompts ...string) error {
|
||||||
|
client, err := NewAPIClient()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, prompt := range prompts {
|
||||||
|
client.Generate(context.Background(), &api.GenerateRequest{Model: model, Prompt: prompt}, func(bts []byte) {
|
||||||
|
var resp api.GenerateResponse
|
||||||
|
if err := json.Unmarshal(bts, &resp); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
fmt.Print(resp.Response)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fmt.Println()
|
||||||
|
fmt.Println()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func serve() error {
|
func generateInteractive(model string) error {
|
||||||
|
fmt.Print(">>> ")
|
||||||
|
scanner := bufio.NewScanner(os.Stdin)
|
||||||
|
for scanner.Scan() {
|
||||||
|
if err := generate(model, scanner.Text()); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
fmt.Print(">>> ")
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func generateBatch(model string) error {
|
||||||
|
scanner := bufio.NewScanner(os.Stdin)
|
||||||
|
for scanner.Scan() {
|
||||||
|
prompt := scanner.Text()
|
||||||
|
fmt.Printf(">>> %s\n", prompt)
|
||||||
|
if err := generate(model, prompt); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func RunServer(_ *cobra.Command, _ []string) error {
|
||||||
ln, err := net.Listen("tcp", "127.0.0.1:11434")
|
ln, err := net.Listen("tcp", "127.0.0.1:11434")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
|
@ -82,39 +147,32 @@ func NewCLI() *cobra.Command {
|
||||||
log.SetFlags(log.LstdFlags | log.Lshortfile)
|
log.SetFlags(log.LstdFlags | log.Lshortfile)
|
||||||
|
|
||||||
rootCmd := &cobra.Command{
|
rootCmd := &cobra.Command{
|
||||||
Use: "ollama",
|
Use: "ollama",
|
||||||
Short: "Large language model runner",
|
Short: "Large language model runner",
|
||||||
|
SilenceUsage: true,
|
||||||
CompletionOptions: cobra.CompletionOptions{
|
CompletionOptions: cobra.CompletionOptions{
|
||||||
DisableDefaultCmd: true,
|
DisableDefaultCmd: true,
|
||||||
},
|
},
|
||||||
PersistentPreRun: func(cmd *cobra.Command, args []string) {
|
PersistentPreRunE: func(_ *cobra.Command, args []string) error {
|
||||||
// Disable usage printing on errors
|
|
||||||
cmd.SilenceUsage = true
|
|
||||||
// create the models directory and it's parent
|
// create the models directory and it's parent
|
||||||
if err := os.MkdirAll(path.Join(cacheDir(), "models"), 0o700); err != nil {
|
return os.MkdirAll(path.Join(cacheDir(), "models"), 0o700)
|
||||||
panic(err)
|
|
||||||
}
|
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
cobra.EnableCommandSorting = false
|
cobra.EnableCommandSorting = false
|
||||||
|
|
||||||
runCmd := &cobra.Command{
|
runCmd := &cobra.Command{
|
||||||
Use: "run MODEL",
|
Use: "run MODEL [PROMPT]",
|
||||||
Short: "Run a model",
|
Short: "Run a model",
|
||||||
Args: cobra.ExactArgs(1),
|
Args: cobra.MinimumNArgs(1),
|
||||||
RunE: func(cmd *cobra.Command, args []string) error {
|
RunE: RunRun,
|
||||||
return run(args[0])
|
|
||||||
},
|
|
||||||
}
|
}
|
||||||
|
|
||||||
serveCmd := &cobra.Command{
|
serveCmd := &cobra.Command{
|
||||||
Use: "serve",
|
Use: "serve",
|
||||||
Aliases: []string{"start"},
|
Aliases: []string{"start"},
|
||||||
Short: "Start ollama",
|
Short: "Start ollama",
|
||||||
RunE: func(cmd *cobra.Command, args []string) error {
|
RunE: RunServer,
|
||||||
return serve()
|
|
||||||
},
|
|
||||||
}
|
}
|
||||||
|
|
||||||
rootCmd.AddCommand(
|
rootCmd.AddCommand(
|
||||||
|
|
1
go.mod
1
go.mod
|
@ -35,6 +35,7 @@ require (
|
||||||
golang.org/x/crypto v0.10.0 // indirect
|
golang.org/x/crypto v0.10.0 // indirect
|
||||||
golang.org/x/net v0.10.0 // indirect
|
golang.org/x/net v0.10.0 // indirect
|
||||||
golang.org/x/sys v0.10.0 // indirect
|
golang.org/x/sys v0.10.0 // indirect
|
||||||
|
golang.org/x/term v0.10.0
|
||||||
golang.org/x/text v0.10.0 // indirect
|
golang.org/x/text v0.10.0 // indirect
|
||||||
google.golang.org/protobuf v1.30.0 // indirect
|
google.golang.org/protobuf v1.30.0 // indirect
|
||||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||||
|
|
2
go.sum
2
go.sum
|
@ -106,6 +106,8 @@ golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||||
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
|
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
|
||||||
|
golang.org/x/term v0.10.0 h1:3R7pNqamzBraeqj/Tj8qt1aQ2HpmlC+Cx/qL/7hn4/c=
|
||||||
|
golang.org/x/term v0.10.0/go.mod h1:lpqdcUyK/oCiQxvxVrppt5ggO2KCZ5QblwqPnfZ6d5o=
|
||||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||||
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
||||||
|
|
Loading…
Reference in a new issue