ollama/cmd/gguf/gguf.go
Blake Mizerany c49947dcf5 init
2024-03-31 09:24:53 -07:00

98 lines
1.9 KiB
Go

// Gguf is a tool for learning about GGUF files.
//
// Usage:
//
// gguf [flags] <file>
package main
import (
"flag"
"fmt"
"io"
"log"
"os"
"text/tabwriter"
"bllamo.com/encoding/gguf"
)
func main() {
if err := Main(os.Stdout, os.Args[1:]...); err != nil {
log.Fatal(err)
}
}
func Main(stdout io.Writer, args ...string) error {
fs := flag.NewFlagSet("gguf", flag.ExitOnError)
flagGPU := fs.Uint64("gpu", 0, "use N bytes of GPU memory (default is 0)")
fs.Usage = func() {
io.WriteString(stdout, "Gguf is a tool for learning about GGUF files.\n")
io.WriteString(stdout, "\n")
io.WriteString(stdout, "Usage:\n")
io.WriteString(stdout, "\n")
io.WriteString(stdout, "\tgguf [flags] <file>\n")
io.WriteString(stdout, "\n")
var numFlags int
fs.VisitAll(func(*flag.Flag) { numFlags++ })
if numFlags > 0 {
io.WriteString(stdout, "Flags:\n")
fs.PrintDefaults()
}
}
fs.Parse(args)
if fs.NArg() != 1 {
fs.Usage()
os.Exit(2)
}
file := fs.Arg(0)
f, err := os.Open(file)
if err != nil {
log.Fatal(err)
}
defer f.Close()
g, err := gguf.ReadFile(f)
if err != nil {
log.Fatal(err)
}
tw := tabwriter.NewWriter(stdout, 0, 2, 2, ' ', 0)
defer tw.Flush()
fmt.Fprintf(tw, "version:\t%d\n", g.Version())
for m, err := range g.Metadata {
if err != nil {
log.Fatal(err)
}
if len(m.Values) > 5 {
fmt.Fprintf(tw, "meta:\t%q: ... (%d values)\n", m.Key, len(m.Values))
} else {
fmt.Fprintf(tw, "meta:\t%q: %v\n", m.Key, m.Values)
}
}
var i int
var totalLayerBytes uint64
var offGPU bool
for t, err := range g.Tensors {
if err != nil {
log.Fatal(err)
}
totalLayerBytes += t.Size
if totalLayerBytes > *flagGPU {
offGPU = true
}
const msg = "tensor (layer %000d):\t%q\t%s\tdims=%v\toffset=%d\tsize=%d\tonGPU=%v\n"
fmt.Fprintf(tw, msg, i, t.Name, t.Type, t.Dimensions, t.Offset, t.Size, !offGPU)
i++
}
return nil
}