49 lines
1.1 KiB
Go
49 lines
1.1 KiB
Go
package main
|
||
|
||
import (
|
||
"bufio"
|
||
"fmt"
|
||
"os"
|
||
"strings"
|
||
|
||
llama "go-llama-new.cpp"
|
||
)
|
||
|
||
func main() {
|
||
if len(os.Args) < 2 {
|
||
fmt.Fprintf(os.Stderr, "usage: %s <model.gguf> [prompt]\n", os.Args[0])
|
||
os.Exit(1)
|
||
}
|
||
|
||
modelPath := os.Args[1]
|
||
prompt := "Hello"
|
||
if len(os.Args) > 2 {
|
||
prompt = strings.Join(os.Args[2:], " ")
|
||
}
|
||
|
||
l, err := llama.New(modelPath, llama.SetContext(512), llama.SetGPULayers(0))
|
||
if err != nil {
|
||
fmt.Fprintf(os.Stderr, "load model: %v\n", err)
|
||
os.Exit(1)
|
||
}
|
||
defer l.Free()
|
||
|
||
out, err := l.Predict(prompt, llama.SetTokens(64), llama.SetThreads(4))
|
||
if err != nil {
|
||
fmt.Fprintf(os.Stderr, "predict: %v\n", err)
|
||
os.Exit(1)
|
||
}
|
||
|
||
fmt.Println(out)
|
||
|
||
reader := bufio.NewReader(os.Stdin)
|
||
fmt.Print("\nТокенизация (введите текст): ")
|
||
line, _ := reader.ReadString('\n')
|
||
_, tokens, err := l.TokenizeString(strings.TrimSpace(line))
|
||
if err != nil {
|
||
fmt.Fprintf(os.Stderr, "tokenize: %v\n", err)
|
||
return
|
||
}
|
||
fmt.Printf("токенов: %d, ids: %v\n", len(tokens), tokens)
|
||
}
|