63 lines
1.1 KiB
Go
63 lines
1.1 KiB
Go
//go:build ai
|
|
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"path"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestAILocal(t *testing.T) {
|
|
ctx, can := context.WithTimeout(context.Background(), time.Minute)
|
|
defer can()
|
|
|
|
d := os.TempDir()
|
|
for k, u := range map[string]string{
|
|
"checkpoints": "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.bin",
|
|
"tokenizer": "https://github.com/karpathy/llama2.c/raw/master/tokenizer.bin",
|
|
} {
|
|
func() {
|
|
if _, err := os.Stat(path.Join(d, k)); os.IsNotExist(err) {
|
|
t.Logf("downloading %s from %s", u, k)
|
|
|
|
resp, err := http.Get(u)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
f, err := os.Create(path.Join(d, k))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer f.Close()
|
|
|
|
if _, err := io.Copy(f, resp.Body); err != nil {
|
|
f.Close()
|
|
os.Remove(path.Join(d, k))
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
ai := NewAILocal(
|
|
path.Join(d, "checkpoints"),
|
|
path.Join(d, "tokenizer"),
|
|
0.9,
|
|
256,
|
|
0.9,
|
|
)
|
|
|
|
if result, err := ai.Do(ctx, "hello world"); err != nil {
|
|
t.Fatal(err)
|
|
} else {
|
|
t.Logf("%s", result)
|
|
}
|
|
}
|