From ecab6f1cc582a5ce8ee2bfbc780cb9990115a3da Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Wed, 11 Sep 2024 11:01:30 -0700 Subject: [PATCH] refactor show ouput fixes line wrapping on long texts --- cmd/cmd.go | 174 +++++++++++++++----------------------- cmd/cmd_test.go | 206 +++++++++++++++++++++++++++++++++++++++++++++ cmd/interactive.go | 2 +- 3 files changed, 277 insertions(+), 105 deletions(-) create mode 100644 cmd/cmd_test.go diff --git a/cmd/cmd.go b/cmd/cmd.go index 5de1ed1b..1fb721e7 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -2,6 +2,7 @@ package cmd import ( "archive/zip" + "bufio" "bytes" "context" "crypto/ed25519" @@ -21,6 +22,7 @@ import ( "regexp" "runtime" "slices" + "strconv" "strings" "sync/atomic" "syscall" @@ -578,7 +580,7 @@ func ListHandler(cmd *cobra.Command, args []string) error { table.SetHeaderLine(false) table.SetBorder(false) table.SetNoWhiteSpace(true) - table.SetTablePadding("\t") + table.SetTablePadding(" ") table.AppendBulk(data) table.Render() @@ -624,7 +626,7 @@ func ListRunningHandler(cmd *cobra.Command, args []string) error { table.SetHeaderLine(false) table.SetBorder(false) table.SetNoWhiteSpace(true) - table.SetTablePadding("\t") + table.SetTablePadding(" ") table.AppendBulk(data) table.Render() @@ -720,125 +722,89 @@ func ShowHandler(cmd *cobra.Command, args []string) error { return nil } - showInfo(resp) - - return nil + return showInfo(resp, os.Stdout) } -func showInfo(resp *api.ShowResponse) { - modelData := [][]string{ - {"parameters", resp.Details.ParameterSize}, - {"quantization", resp.Details.QuantizationLevel}, - } - if resp.ModelInfo != nil { - arch := resp.ModelInfo["general.architecture"].(string) - modelData = append(modelData, - []string{"arch", arch}, - []string{"context length", fmt.Sprintf("%v", resp.ModelInfo[fmt.Sprintf("%s.context_length", arch)].(float64))}, - []string{"embedding length", fmt.Sprintf("%v", resp.ModelInfo[fmt.Sprintf("%s.embedding_length", arch)].(float64))}, - ) +func showInfo(resp *api.ShowResponse, w io.Writer) error { + tableRender := func(header string, rows func() [][]string) { + fmt.Fprintln(w, " ", header) + table := tablewriter.NewWriter(w) + table.SetAlignment(tablewriter.ALIGN_LEFT) + table.SetBorder(false) + table.SetNoWhiteSpace(true) + table.SetTablePadding(" ") + + switch header { + case "Template", "System", "License": + table.SetColWidth(100) + } + + table.AppendBulk(rows()) + table.Render() + fmt.Fprintln(w) } - mainTableData := [][]string{ - {"Model"}, - {renderSubTable(modelData, false)}, - } + tableRender("Model", func() (rows [][]string) { + if resp.ModelInfo != nil { + arch := resp.ModelInfo["general.architecture"].(string) + rows = append(rows, []string{"", "architecture", arch}) + rows = append(rows, []string{"", "parameters", format.HumanNumber(uint64(resp.ModelInfo["general.parameter_count"].(float64)))}) + rows = append(rows, []string{"", "context length", strconv.FormatFloat(resp.ModelInfo[fmt.Sprintf("%s.context_length", arch)].(float64), 'f', -1, 64)}) + rows = append(rows, []string{"", "embedding length", strconv.FormatFloat(resp.ModelInfo[fmt.Sprintf("%s.embedding_length", arch)].(float64), 'f', -1, 64)}) + } else { + rows = append(rows, []string{"", "architecture", resp.Details.Family}) + rows = append(rows, []string{"", "parameters", resp.Details.ParameterSize}) + } + rows = append(rows, []string{"", "quantization", resp.Details.QuantizationLevel}) + return + }) if resp.ProjectorInfo != nil { - projectorData := [][]string{ - {"arch", "clip"}, - {"parameters", format.HumanNumber(uint64(resp.ProjectorInfo["general.parameter_count"].(float64)))}, - } - - if projectorType, ok := resp.ProjectorInfo["clip.projector_type"]; ok { - projectorData = append(projectorData, []string{"projector type", projectorType.(string)}) - } - - projectorData = append(projectorData, - []string{"embedding length", fmt.Sprintf("%v", resp.ProjectorInfo["clip.vision.embedding_length"].(float64))}, - []string{"projection dimensionality", fmt.Sprintf("%v", resp.ProjectorInfo["clip.vision.projection_dim"].(float64))}, - ) - - mainTableData = append(mainTableData, - []string{"Projector"}, - []string{renderSubTable(projectorData, false)}, - ) + tableRender("Projector", func() (rows [][]string) { + arch := resp.ProjectorInfo["general.architecture"].(string) + rows = append(rows, []string{"", "architecture", arch}) + rows = append(rows, []string{"", "parameters", format.HumanNumber(uint64(resp.ProjectorInfo["general.parameter_count"].(float64)))}) + rows = append(rows, []string{"", "embedding length", strconv.FormatFloat(resp.ProjectorInfo[fmt.Sprintf("%s.vision.embedding_length", arch)].(float64), 'f', -1, 64)}) + rows = append(rows, []string{"", "dimensions", strconv.FormatFloat(resp.ProjectorInfo[fmt.Sprintf("%s.vision.projection_dim", arch)].(float64), 'f', -1, 64)}) + return + }) } if resp.Parameters != "" { - mainTableData = append(mainTableData, []string{"Parameters"}, []string{formatParams(resp.Parameters)}) + tableRender("Parameters", func() (rows [][]string) { + scanner := bufio.NewScanner(strings.NewReader(resp.Parameters)) + for scanner.Scan() { + if text := scanner.Text(); text != "" { + rows = append(rows, append([]string{""}, strings.Fields(text)...)) + } + } + return + }) + } + + head := func(s string, n int) (rows [][]string) { + scanner := bufio.NewScanner(strings.NewReader(s)) + for scanner.Scan() && (len(rows) < n || n < 0) { + if text := scanner.Text(); text != "" { + rows = append(rows, []string{"", strings.TrimSpace(text)}) + } + } + return } if resp.System != "" { - mainTableData = append(mainTableData, []string{"System"}, []string{renderSubTable(twoLines(resp.System), true)}) + tableRender("System", func() [][]string { + return head(resp.System, 2) + }) } if resp.License != "" { - mainTableData = append(mainTableData, []string{"License"}, []string{renderSubTable(twoLines(resp.License), true)}) + tableRender("License", func() [][]string { + return head(resp.License, 2) + }) } - table := tablewriter.NewWriter(os.Stdout) - table.SetAutoWrapText(false) - table.SetBorder(false) - table.SetAlignment(tablewriter.ALIGN_LEFT) - - for _, v := range mainTableData { - table.Append(v) - } - - table.Render() -} - -func renderSubTable(data [][]string, file bool) string { - var buf bytes.Buffer - table := tablewriter.NewWriter(&buf) - table.SetAutoWrapText(!file) - table.SetBorder(false) - table.SetNoWhiteSpace(true) - table.SetTablePadding("\t") - table.SetAlignment(tablewriter.ALIGN_LEFT) - - for _, v := range data { - table.Append(v) - } - - table.Render() - - renderedTable := buf.String() - lines := strings.Split(renderedTable, "\n") - for i, line := range lines { - lines[i] = "\t" + line - } - - return strings.Join(lines, "\n") -} - -func twoLines(s string) [][]string { - lines := strings.Split(s, "\n") - res := [][]string{} - - count := 0 - for _, line := range lines { - line = strings.TrimSpace(line) - if line != "" { - count++ - res = append(res, []string{line}) - if count == 2 { - return res - } - } - } - return res -} - -func formatParams(s string) string { - lines := strings.Split(s, "\n") - table := [][]string{} - - for _, line := range lines { - table = append(table, strings.Fields(line)) - } - return renderSubTable(table, false) + return nil } func CopyHandler(cmd *cobra.Command, args []string) error { diff --git a/cmd/cmd_test.go b/cmd/cmd_test.go new file mode 100644 index 00000000..0f8863cc --- /dev/null +++ b/cmd/cmd_test.go @@ -0,0 +1,206 @@ +package cmd + +import ( + "bytes" + "os" + "path/filepath" + "testing" + + "github.com/google/go-cmp/cmp" + + "github.com/ollama/ollama/api" +) + +func TestShowInfo(t *testing.T) { + t.Run("bare details", func(t *testing.T) { + var b bytes.Buffer + if err := showInfo(&api.ShowResponse{ + Details: api.ModelDetails{ + Family: "test", + ParameterSize: "7B", + QuantizationLevel: "FP16", + }, + }, &b); err != nil { + t.Fatal(err) + } + + expect := ` Model + architecture test + parameters 7B + quantization FP16 + +` + + if diff := cmp.Diff(expect, b.String()); diff != "" { + t.Errorf("unexpected output (-want +got):\n%s", diff) + } + }) + + t.Run("bare model info", func(t *testing.T) { + var b bytes.Buffer + if err := showInfo(&api.ShowResponse{ + ModelInfo: map[string]any{ + "general.architecture": "test", + "general.parameter_count": float64(7_000_000_000), + "test.context_length": float64(0), + "test.embedding_length": float64(0), + }, + Details: api.ModelDetails{ + Family: "test", + ParameterSize: "7B", + QuantizationLevel: "FP16", + }, + }, &b); err != nil { + t.Fatal(err) + } + + expect := ` Model + architecture test + parameters 7B + context length 0 + embedding length 0 + quantization FP16 + +` + if diff := cmp.Diff(expect, b.String()); diff != "" { + t.Errorf("unexpected output (-want +got):\n%s", diff) + } + }) + + t.Run("parameters", func(t *testing.T) { + var b bytes.Buffer + if err := showInfo(&api.ShowResponse{ + Details: api.ModelDetails{ + Family: "test", + ParameterSize: "7B", + QuantizationLevel: "FP16", + }, + Parameters: ` + stop never + stop gonna + stop give + stop you + stop up + temperature 99`, + }, &b); err != nil { + t.Fatal(err) + } + + expect := ` Model + architecture test + parameters 7B + quantization FP16 + + Parameters + stop never + stop gonna + stop give + stop you + stop up + temperature 99 + +` + if diff := cmp.Diff(expect, b.String()); diff != "" { + t.Errorf("unexpected output (-want +got):\n%s", diff) + } + }) + + t.Run("project info", func(t *testing.T) { + var b bytes.Buffer + if err := showInfo(&api.ShowResponse{ + Details: api.ModelDetails{ + Family: "test", + ParameterSize: "7B", + QuantizationLevel: "FP16", + }, + ProjectorInfo: map[string]any{ + "general.architecture": "clip", + "general.parameter_count": float64(133_700_000), + "clip.vision.embedding_length": float64(0), + "clip.vision.projection_dim": float64(0), + }, + }, &b); err != nil { + t.Fatal(err) + } + + expect := ` Model + architecture test + parameters 7B + quantization FP16 + + Projector + architecture clip + parameters 133.70M + embedding length 0 + dimensions 0 + +` + if diff := cmp.Diff(expect, b.String()); diff != "" { + t.Errorf("unexpected output (-want +got):\n%s", diff) + } + }) + + t.Run("system", func(t *testing.T) { + var b bytes.Buffer + if err := showInfo(&api.ShowResponse{ + Details: api.ModelDetails{ + Family: "test", + ParameterSize: "7B", + QuantizationLevel: "FP16", + }, + System: `You are a pirate! +Ahoy, matey! +Weigh anchor! + `, + }, &b); err != nil { + t.Fatal(err) + } + + expect := ` Model + architecture test + parameters 7B + quantization FP16 + + System + You are a pirate! + Ahoy, matey! + +` + if diff := cmp.Diff(expect, b.String()); diff != "" { + t.Errorf("unexpected output (-want +got):\n%s", diff) + } + }) + + t.Run("license", func(t *testing.T) { + var b bytes.Buffer + license, err := os.ReadFile(filepath.Join("..", "LICENSE")) + if err != nil { + t.Fatal(err) + } + + if err := showInfo(&api.ShowResponse{ + Details: api.ModelDetails{ + Family: "test", + ParameterSize: "7B", + QuantizationLevel: "FP16", + }, + License: string(license), + }, &b); err != nil { + t.Fatal(err) + } + + expect := ` Model + architecture test + parameters 7B + quantization FP16 + + License + MIT License + Copyright (c) Ollama + +` + if diff := cmp.Diff(expect, b.String()); diff != "" { + t.Errorf("unexpected output (-want +got):\n%s", diff) + } + }) +} diff --git a/cmd/interactive.go b/cmd/interactive.go index 4462cf29..9fe1ed4c 100644 --- a/cmd/interactive.go +++ b/cmd/interactive.go @@ -371,7 +371,7 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error { switch args[1] { case "info": - showInfo(resp) + _ = showInfo(resp, os.Stderr) case "license": if resp.License == "" { fmt.Println("No license was specified for this model.")