diff options
Diffstat (limited to 'research/sander/extract.go')
| -rw-r--r-- | research/sander/extract.go | 11 |
1 files changed, 6 insertions, 5 deletions
diff --git a/research/sander/extract.go b/research/sander/extract.go index de31b1f..c72fe7f 100644 --- a/research/sander/extract.go +++ b/research/sander/extract.go @@ -3,6 +3,7 @@ package sander import ( "database/sql" "database/sql/driver" + "fmt" "go.jknobloc.com/x/onnx" "go.jknobloc.com/x/research/lesci" @@ -20,16 +21,16 @@ func (e *Experiment) Extract(db *sql.DB) error { shape = s } - if shape[1] != 768 { - panic("unimplemented") + if shape[1] != e.hiddenDim { + panic("shape mismatch") } t := tensor.NewDense[float32](shape, data) - if _, err := db.Exec(` + if _, err := db.Exec(fmt.Sprintf(` DROP TABLE IF EXISTS embeddings; - CREATE TABLE embeddings (token_id INTEGER, embedding FLOAT[768], reference BOOLEAN); - `); err != nil { + CREATE TABLE embeddings (token_id INTEGER, embedding FLOAT[%d], reference BOOLEAN); + `, e.hiddenDim)); err != nil { return err } |