summaryrefslogtreecommitdiff
path: root/gpt2/cmd
diff options
context:
space:
mode:
Diffstat (limited to 'gpt2/cmd')
-rw-r--r--gpt2/cmd/gpt2/main.go5
1 files changed, 3 insertions, 2 deletions
diff --git a/gpt2/cmd/gpt2/main.go b/gpt2/cmd/gpt2/main.go
index 06ec535..b9b75a2 100644
--- a/gpt2/cmd/gpt2/main.go
+++ b/gpt2/cmd/gpt2/main.go
@@ -7,6 +7,7 @@ import (
"slices"
"go.jknobloc.com/x/gpt2"
+ "go.jknobloc.com/x/shelf"
)
func main() {
@@ -37,7 +38,7 @@ func generate(prompt []int64) {
WithLogProbs: false,
}
- m := gpt2.NewModel("gpt2/models/mbpe_conv/gpt2_8192_m000_babylm_v2/model_cache.onnx", cfg, opts)
+ m := gpt2.NewModel(shelf.Abs("models/mbpe/gpt2_8192_m000_babylm_v2/model_cache.onnx"), cfg, opts)
if err := m.Init(); err != nil {
log.Fatal(err)
@@ -69,7 +70,7 @@ func score(prompt []int64) {
WithLogProbs: true,
}
- m := gpt2.NewModel("gpt2/models/mbpe_conv/gpt2_8192_m000_babylm_v2/model_eval.onnx", cfg, opts)
+ m := gpt2.NewModel(shelf.Abs("models/mbpe/gpt2_8192_m000_babylm_v2/model_eval.onnx"), cfg, opts)
if err := m.Init(); err != nil {
log.Fatal(err)