2023-07-03 19:22:44 +00:00
|
|
|
package cmd
|
|
|
|
|
|
|
|
import (
|
2023-07-06 20:49:31 +00:00
|
|
|
"bufio"
|
2023-07-03 19:22:44 +00:00
|
|
|
"context"
|
2023-07-06 22:43:04 +00:00
|
|
|
"errors"
|
2023-07-06 16:24:49 +00:00
|
|
|
"fmt"
|
2023-07-03 19:22:44 +00:00
|
|
|
"log"
|
|
|
|
"net"
|
2023-07-11 20:05:51 +00:00
|
|
|
"net/http"
|
2023-07-03 19:22:44 +00:00
|
|
|
"os"
|
2023-07-15 00:27:14 +00:00
|
|
|
"path/filepath"
|
2023-07-07 17:12:58 +00:00
|
|
|
"strings"
|
2023-07-06 22:43:04 +00:00
|
|
|
"time"
|
2023-07-03 19:22:44 +00:00
|
|
|
|
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-03 19:22:44 +00:00
|
|
|
)
|
|
|
|
|
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-15 00:27:14 +00:00
|
|
|
return filepath.Join(home, ".ollama")
|
2023-07-04 04:47:00 +00:00
|
|
|
}
|
|
|
|
|
2023-07-17 00:02:22 +00:00
|
|
|
func create(cmd *cobra.Command, args []string) error {
|
|
|
|
filename, _ := cmd.Flags().GetString("file")
|
|
|
|
client := api.NewClient()
|
|
|
|
|
|
|
|
request := api.CreateRequest{Name: args[0], Path: filename}
|
|
|
|
fn := func(resp api.CreateProgress) error {
|
|
|
|
fmt.Println(resp.Status)
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
if err := client.Create(context.Background(), &request, fn); err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
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-17 00:02:22 +00:00
|
|
|
func push(cmd *cobra.Command, args []string) error {
|
|
|
|
client := api.NewClient()
|
|
|
|
|
|
|
|
request := api.PushRequest{Name: args[0]}
|
|
|
|
fn := func(resp api.PushProgress) error {
|
|
|
|
fmt.Println(resp.Status)
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
if err := client.Push(context.Background(), &request, fn); err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func RunPull(cmd *cobra.Command, args []string) error {
|
|
|
|
return pull(args[0])
|
|
|
|
}
|
|
|
|
|
2023-07-06 21:05:55 +00:00
|
|
|
func pull(model string) error {
|
2023-07-07 18:26:58 +00:00
|
|
|
client := api.NewClient()
|
2023-07-17 00:02:22 +00:00
|
|
|
|
2023-07-07 18:26:58 +00:00
|
|
|
var bar *progressbar.ProgressBar
|
2023-07-11 20:05:51 +00:00
|
|
|
|
2023-07-17 00:02:22 +00:00
|
|
|
currentLayer := ""
|
|
|
|
request := api.PullRequest{Name: model}
|
|
|
|
fn := func(resp api.PullProgress) error {
|
|
|
|
if resp.Digest != currentLayer && resp.Digest != "" {
|
|
|
|
if currentLayer != "" {
|
|
|
|
fmt.Println()
|
2023-07-07 18:26:58 +00:00
|
|
|
}
|
2023-07-17 00:02:22 +00:00
|
|
|
currentLayer = resp.Digest
|
|
|
|
layerStr := resp.Digest[7:23] + "..."
|
|
|
|
bar = progressbar.DefaultBytes(
|
|
|
|
int64(resp.Total),
|
|
|
|
"pulling "+layerStr,
|
|
|
|
)
|
|
|
|
} else if resp.Digest == currentLayer && resp.Digest != "" {
|
|
|
|
bar.Set(resp.Completed)
|
|
|
|
} else {
|
|
|
|
currentLayer = ""
|
|
|
|
fmt.Println(resp.Status)
|
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|
2023-07-07 14:22:37 +00:00
|
|
|
|
2023-07-17 00:02:22 +00:00
|
|
|
if err := client.Pull(context.Background(), &request, fn); err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
return nil
|
2023-07-06 20:49:31 +00:00
|
|
|
}
|
|
|
|
|
2023-07-13 01:18:06 +00:00
|
|
|
func RunGenerate(cmd *cobra.Command, args []string) error {
|
2023-07-06 20:49:31 +00:00
|
|
|
if len(args) > 1 {
|
2023-07-12 16:55:07 +00:00
|
|
|
// join all args into a single prompt
|
2023-07-13 01:18:06 +00:00
|
|
|
return generate(cmd, args[0], strings.Join(args[1:], " "))
|
2023-07-06 20:49:31 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
if term.IsTerminal(int(os.Stdin.Fd())) {
|
2023-07-13 01:18:06 +00:00
|
|
|
return generateInteractive(cmd, args[0])
|
2023-07-06 20:49:31 +00:00
|
|
|
}
|
|
|
|
|
2023-07-13 01:18:06 +00:00
|
|
|
return generateBatch(cmd, args[0])
|
2023-07-06 20:49:31 +00:00
|
|
|
}
|
|
|
|
|
2023-07-13 18:02:53 +00:00
|
|
|
var generateContextKey struct{}
|
|
|
|
|
2023-07-13 01:18:06 +00:00
|
|
|
func generate(cmd *cobra.Command, 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-13 01:18:06 +00:00
|
|
|
var latest api.GenerateResponse
|
|
|
|
|
2023-07-13 18:02:53 +00:00
|
|
|
generateContext, ok := cmd.Context().Value(generateContextKey).([]int)
|
|
|
|
if !ok {
|
|
|
|
generateContext = []int{}
|
|
|
|
}
|
|
|
|
|
|
|
|
request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext}
|
2023-07-07 21:04:43 +00:00
|
|
|
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-13 01:18:06 +00:00
|
|
|
latest = resp
|
|
|
|
|
2023-07-07 17:12:58 +00:00
|
|
|
fmt.Print(resp.Response)
|
2023-07-13 18:02:53 +00:00
|
|
|
|
|
|
|
cmd.SetContext(context.WithValue(cmd.Context(), generateContextKey, resp.Context))
|
2023-07-07 17:12:58 +00:00
|
|
|
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-13 01:18:06 +00:00
|
|
|
|
|
|
|
verbose, err := cmd.Flags().GetBool("verbose")
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
if verbose {
|
|
|
|
latest.Summary()
|
|
|
|
}
|
2023-07-07 17:12:58 +00:00
|
|
|
}
|
2023-07-06 20:49:31 +00:00
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2023-07-13 01:18:06 +00:00
|
|
|
func generateInteractive(cmd *cobra.Command, model string) error {
|
2023-07-06 20:49:31 +00:00
|
|
|
fmt.Print(">>> ")
|
|
|
|
scanner := bufio.NewScanner(os.Stdin)
|
|
|
|
for scanner.Scan() {
|
2023-07-13 01:18:06 +00:00
|
|
|
if err := generate(cmd, model, scanner.Text()); err != nil {
|
2023-07-06 20:49:31 +00:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
fmt.Print(">>> ")
|
|
|
|
}
|
|
|
|
|
2023-07-06 18:57:11 +00:00
|
|
|
return nil
|
2023-07-06 16:24:49 +00:00
|
|
|
}
|
|
|
|
|
2023-07-13 01:18:06 +00:00
|
|
|
func generateBatch(cmd *cobra.Command, model string) error {
|
2023-07-06 20:49:31 +00:00
|
|
|
scanner := bufio.NewScanner(os.Stdin)
|
|
|
|
for scanner.Scan() {
|
|
|
|
prompt := scanner.Text()
|
|
|
|
fmt.Printf(">>> %s\n", prompt)
|
2023-07-13 01:18:06 +00:00
|
|
|
if err := generate(cmd, model, prompt); err != nil {
|
2023-07-06 20:49:31 +00:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func RunServer(_ *cobra.Command, _ []string) error {
|
2023-07-07 20:48:13 +00:00
|
|
|
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)
|
|
|
|
}
|
|
|
|
|
2023-07-03 19:22:44 +00:00
|
|
|
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,
|
2023-07-03 19:22:44 +00:00
|
|
|
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-15 00:27:14 +00:00
|
|
|
return os.MkdirAll(filepath.Join(cacheDir(), "models"), 0o700)
|
2023-07-03 19:22:44 +00:00
|
|
|
},
|
|
|
|
}
|
|
|
|
|
|
|
|
cobra.EnableCommandSorting = false
|
|
|
|
|
2023-07-17 00:02:22 +00:00
|
|
|
createCmd := &cobra.Command{
|
|
|
|
Use: "create MODEL",
|
|
|
|
Short: "Create a model from a Modelfile",
|
|
|
|
Args: cobra.MinimumNArgs(1),
|
|
|
|
RunE: create,
|
|
|
|
}
|
|
|
|
|
|
|
|
createCmd.Flags().StringP("file", "f", "Modelfile", "Name of the Modelfile (default \"Modelfile\")")
|
|
|
|
|
2023-07-03 19:22:44 +00:00
|
|
|
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,
|
2023-07-03 19:22:44 +00:00
|
|
|
}
|
|
|
|
|
2023-07-13 01:18:06 +00:00
|
|
|
runCmd.Flags().Bool("verbose", false, "Show timings for response")
|
|
|
|
|
2023-07-03 19:22:44 +00:00
|
|
|
serveCmd := &cobra.Command{
|
|
|
|
Use: "serve",
|
|
|
|
Aliases: []string{"start"},
|
|
|
|
Short: "Start ollama",
|
2023-07-06 20:49:31 +00:00
|
|
|
RunE: RunServer,
|
2023-07-03 19:22:44 +00:00
|
|
|
}
|
|
|
|
|
2023-07-17 00:02:22 +00:00
|
|
|
pullCmd := &cobra.Command{
|
|
|
|
Use: "pull MODEL",
|
|
|
|
Short: "Pull a model from a registry",
|
|
|
|
Args: cobra.MinimumNArgs(1),
|
|
|
|
RunE: RunPull,
|
|
|
|
}
|
|
|
|
|
|
|
|
pushCmd := &cobra.Command{
|
|
|
|
Use: "push MODEL",
|
|
|
|
Short: "Push a model to a registry",
|
|
|
|
Args: cobra.MinimumNArgs(1),
|
|
|
|
RunE: push,
|
|
|
|
}
|
|
|
|
|
2023-07-03 19:22:44 +00:00
|
|
|
rootCmd.AddCommand(
|
|
|
|
serveCmd,
|
2023-07-17 00:02:22 +00:00
|
|
|
createCmd,
|
2023-07-03 21:14:20 +00:00
|
|
|
runCmd,
|
2023-07-17 00:02:22 +00:00
|
|
|
pullCmd,
|
|
|
|
pushCmd,
|
2023-07-03 19:22:44 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
return rootCmd
|
|
|
|
}
|