ollama/cmd/cmd.go

217 lines
4 KiB
Go
Raw Normal View History

package cmd
import (
2023-07-06 20:49:31 +00:00
"bufio"
"context"
2023-07-06 22:43:04 +00:00
"errors"
2023-07-06 16:24:49 +00:00
"fmt"
"log"
"net"
2023-07-11 20:05:51 +00:00
"net/http"
"os"
"path"
2023-07-07 17:12:58 +00:00
"strings"
2023-07-06 22:43:04 +00:00
"time"
2023-07-06 21:05:55 +00:00
"github.com/schollz/progressbar/v3"
2023-07-06 20:49:31 +00:00
"github.com/spf13/cobra"
"golang.org/x/term"
2023-07-03 20:32:48 +00:00
"github.com/jmorganca/ollama/api"
"github.com/jmorganca/ollama/server"
)
2023-07-06 15:59:42 +00:00
func cacheDir() string {
2023-07-04 04:47:00 +00:00
home, err := os.UserHomeDir()
if err != nil {
panic(err)
}
2023-07-06 15:59:42 +00:00
return path.Join(home, ".ollama")
2023-07-04 04:47:00 +00:00
}
2023-07-06 21:05:55 +00:00
func RunRun(cmd *cobra.Command, args []string) error {
2023-07-06 22:43:04 +00:00
_, err := os.Stat(args[0])
switch {
case errors.Is(err, os.ErrNotExist):
if err := pull(args[0]); err != nil {
2023-07-11 20:05:51 +00:00
var apiStatusError api.StatusError
if !errors.As(err, &apiStatusError) {
return err
}
if apiStatusError.StatusCode != http.StatusBadGateway {
return err
}
2023-07-06 22:43:04 +00:00
}
case err != nil:
2023-07-06 21:05:55 +00:00
return err
}
return RunGenerate(cmd, args)
2023-07-06 18:18:40 +00:00
}
2023-07-06 21:05:55 +00:00
func pull(model string) error {
2023-07-07 18:26:58 +00:00
client := api.NewClient()
var bar *progressbar.ProgressBar
return client.Pull(
context.Background(),
&api.PullRequest{Model: model},
func(progress api.PullProgress) error {
if bar == nil {
2023-07-12 16:55:07 +00:00
if progress.Percent >= 100 {
2023-07-11 20:05:51 +00:00
// already downloaded
return nil
}
2023-07-07 18:26:58 +00:00
bar = progressbar.DefaultBytes(progress.Total)
}
2023-07-07 14:22:37 +00:00
2023-07-07 18:26:58 +00:00
return bar.Set64(progress.Completed)
},
)
2023-07-06 20:49:31 +00:00
}
func RunGenerate(_ *cobra.Command, args []string) error {
if len(args) > 1 {
2023-07-12 16:55:07 +00:00
// join all args into a single prompt
return generate(args[0], strings.Join(args[1:], " "))
2023-07-06 20:49:31 +00:00
}
if term.IsTerminal(int(os.Stdin.Fd())) {
return generateInteractive(args[0])
}
return generateBatch(args[0])
}
2023-07-06 22:43:04 +00:00
func generate(model, prompt string) error {
2023-07-07 17:12:58 +00:00
if len(strings.TrimSpace(prompt)) > 0 {
client := api.NewClient()
spinner := progressbar.NewOptions(-1,
progressbar.OptionSetWriter(os.Stderr),
progressbar.OptionThrottle(60*time.Millisecond),
progressbar.OptionSpinnerType(14),
progressbar.OptionSetRenderBlankState(true),
progressbar.OptionSetElapsedTime(false),
progressbar.OptionClearOnFinish(),
)
go func() {
for range time.Tick(60 * time.Millisecond) {
if spinner.IsFinished() {
break
}
spinner.Add(1)
2023-07-06 22:43:04 +00:00
}
2023-07-07 17:12:58 +00:00
}()
2023-07-06 22:43:04 +00:00
2023-07-07 21:04:43 +00:00
request := api.GenerateRequest{Model: model, Prompt: prompt}
fn := func(resp api.GenerateResponse) error {
2023-07-07 17:12:58 +00:00
if !spinner.IsFinished() {
spinner.Finish()
}
2023-07-06 22:43:04 +00:00
2023-07-07 17:12:58 +00:00
fmt.Print(resp.Response)
return nil
2023-07-07 21:04:43 +00:00
}
if err := client.Generate(context.Background(), &request, fn); err != nil {
return err
}
2023-07-06 22:43:04 +00:00
2023-07-07 17:12:58 +00:00
fmt.Println()
fmt.Println()
}
2023-07-06 20:49:31 +00:00
return nil
}
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(">>> ")
}
2023-07-06 18:57:11 +00:00
return nil
2023-07-06 16:24:49 +00:00
}
2023-07-06 20:49:31 +00:00
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 {
host := os.Getenv("OLLAMA_HOST")
if host == "" {
host = "127.0.0.1"
}
port := os.Getenv("OLLAMA_PORT")
if port == "" {
port = "11434"
}
ln, err := net.Listen("tcp", fmt.Sprintf("%s:%s", host, port))
2023-07-04 04:47:00 +00:00
if err != nil {
return err
}
return server.Serve(ln)
}
func NewCLI() *cobra.Command {
log.SetFlags(log.LstdFlags | log.Lshortfile)
rootCmd := &cobra.Command{
2023-07-06 20:49:31 +00:00
Use: "ollama",
Short: "Large language model runner",
SilenceUsage: true,
CompletionOptions: cobra.CompletionOptions{
DisableDefaultCmd: true,
},
2023-07-06 20:49:31 +00:00
PersistentPreRunE: func(_ *cobra.Command, args []string) error {
2023-07-06 15:59:42 +00:00
// create the models directory and it's parent
2023-07-06 20:49:31 +00:00
return os.MkdirAll(path.Join(cacheDir(), "models"), 0o700)
},
}
cobra.EnableCommandSorting = false
runCmd := &cobra.Command{
2023-07-06 20:49:31 +00:00
Use: "run MODEL [PROMPT]",
2023-07-04 04:47:00 +00:00
Short: "Run a model",
2023-07-06 20:49:31 +00:00
Args: cobra.MinimumNArgs(1),
RunE: RunRun,
}
serveCmd := &cobra.Command{
Use: "serve",
Aliases: []string{"start"},
Short: "Start ollama",
2023-07-06 20:49:31 +00:00
RunE: RunServer,
}
rootCmd.AddCommand(
serveCmd,
runCmd,
)
return rootCmd
}