|
1 | 1 | package pipelineBackends |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "context" |
4 | 5 | "errors" |
| 6 | + "fmt" |
| 7 | + "io" |
| 8 | + "os" |
| 9 | + "strings" |
| 10 | + |
| 11 | + jsoniter "github.com/json-iterator/go" |
5 | 12 |
|
6 | 13 | "github.com/knights-analytics/hugot/options" |
| 14 | + "github.com/knights-analytics/hugot/util" |
7 | 15 | ) |
8 | 16 |
|
9 | 17 | type Model struct { |
10 | | - Path string |
11 | | - OnnxFilename string |
12 | | - OnnxBytes []byte |
13 | | - ORTModel *ORTModel |
14 | | - GoMLXModel *GoMLXModel |
15 | | - Tokenizer *Tokenizer |
16 | | - InputsMeta []InputOutputInfo |
17 | | - OutputsMeta []InputOutputInfo |
18 | | - Destroy func() error |
19 | | - Pipelines map[string]Pipeline |
| 18 | + Path string |
| 19 | + OnnxFilename string |
| 20 | + OnnxBytes []byte |
| 21 | + ORTModel *ORTModel |
| 22 | + GoMLXModel *GoMLXModel |
| 23 | + Tokenizer *Tokenizer |
| 24 | + InputsMeta []InputOutputInfo |
| 25 | + OutputsMeta []InputOutputInfo |
| 26 | + Destroy func() error |
| 27 | + Pipelines map[string]Pipeline |
| 28 | + MaxPositionEmbeddings int |
| 29 | +} |
| 30 | + |
| 31 | +func LoadModel(path string, onnxFilename string, options *options.Options) (*Model, error) { |
| 32 | + model := &Model{ |
| 33 | + Path: path, |
| 34 | + OnnxFilename: onnxFilename, |
| 35 | + Pipelines: make(map[string]Pipeline), |
| 36 | + } |
| 37 | + |
| 38 | + err := LoadOnnxModelBytes(model) |
| 39 | + if err != nil { |
| 40 | + return nil, err |
| 41 | + } |
| 42 | + |
| 43 | + err = loadModelConfig(model) |
| 44 | + if err != nil { |
| 45 | + return nil, err |
| 46 | + } |
| 47 | + |
| 48 | + err = CreateModelBackend(model, options) |
| 49 | + if err != nil { |
| 50 | + return nil, err |
| 51 | + } |
| 52 | + |
| 53 | + tkErr := LoadTokenizer(model, options) |
| 54 | + if tkErr != nil { |
| 55 | + return nil, tkErr |
| 56 | + } |
| 57 | + |
| 58 | + model.Destroy = func() error { |
| 59 | + destroyErr := model.Tokenizer.Destroy() |
| 60 | + switch options.Backend { |
| 61 | + case "ORT": |
| 62 | + destroyErr = errors.Join(destroyErr, model.ORTModel.Destroy()) |
| 63 | + case "GO", "XLA": |
| 64 | + model.GoMLXModel.Destroy() |
| 65 | + } |
| 66 | + return destroyErr |
| 67 | + } |
| 68 | + return model, nil |
| 69 | +} |
| 70 | + |
| 71 | +func LoadOnnxModelBytes(model *Model) error { |
| 72 | + var modelOnnxFile string |
| 73 | + onnxFiles, err := getOnnxFiles(model.Path) |
| 74 | + if err != nil { |
| 75 | + return err |
| 76 | + } |
| 77 | + if len(onnxFiles) == 0 { |
| 78 | + return fmt.Errorf("no .onnx file detected at %s. There should be exactly .onnx file", model.Path) |
| 79 | + } |
| 80 | + if len(onnxFiles) > 1 { |
| 81 | + if model.OnnxFilename == "" { |
| 82 | + return fmt.Errorf("multiple .onnx file detected at %s and no OnnxFilename specified", model.Path) |
| 83 | + } |
| 84 | + modelNameFound := false |
| 85 | + for i := range onnxFiles { |
| 86 | + if onnxFiles[i][1] == model.OnnxFilename { |
| 87 | + modelNameFound = true |
| 88 | + modelOnnxFile = util.PathJoinSafe(onnxFiles[i]...) |
| 89 | + } |
| 90 | + } |
| 91 | + if !modelNameFound { |
| 92 | + return fmt.Errorf("file %s not found at %s", model.OnnxFilename, model.Path) |
| 93 | + } |
| 94 | + } else { |
| 95 | + modelOnnxFile = util.PathJoinSafe(onnxFiles[0]...) |
| 96 | + } |
| 97 | + |
| 98 | + onnxBytes, err := util.ReadFileBytes(modelOnnxFile) |
| 99 | + if err != nil { |
| 100 | + return err |
| 101 | + } |
| 102 | + |
| 103 | + model.OnnxBytes = onnxBytes |
| 104 | + |
| 105 | + return err |
| 106 | +} |
| 107 | + |
| 108 | +func getOnnxFiles(path string) ([][]string, error) { |
| 109 | + var onnxFiles [][]string |
| 110 | + walker := func(_ context.Context, _ string, parent string, info os.FileInfo, _ io.Reader) (toContinue bool, err error) { |
| 111 | + if strings.HasSuffix(info.Name(), ".onnx") { |
| 112 | + onnxFiles = append(onnxFiles, []string{util.PathJoinSafe(path, parent), info.Name()}) |
| 113 | + } |
| 114 | + return true, nil |
| 115 | + } |
| 116 | + err := util.WalkDir()(context.Background(), path, walker) |
| 117 | + return onnxFiles, err |
| 118 | +} |
| 119 | + |
| 120 | +func loadModelConfig(model *Model) error { |
| 121 | + |
| 122 | + // load config.json if it exists, to determine max_position_embeddings |
| 123 | + configPath := util.PathJoinSafe(model.Path, "config.json") |
| 124 | + |
| 125 | + exists, err := util.FileExists(configPath) |
| 126 | + if err != nil { |
| 127 | + return err |
| 128 | + } |
| 129 | + if exists { |
| 130 | + configBytes, readErr := util.ReadFileBytes(configPath) |
| 131 | + if readErr != nil { |
| 132 | + return readErr |
| 133 | + } |
| 134 | + |
| 135 | + configMap := map[string]any{} |
| 136 | + readErr = jsoniter.Unmarshal(configBytes, &configMap) |
| 137 | + if readErr != nil { |
| 138 | + return readErr |
| 139 | + } |
| 140 | + |
| 141 | + if maxPositionEmbeddingsRaw, existsOk := configMap["max_position_embeddings"]; existsOk { |
| 142 | + if maxPositionEmbeddings, castOk := maxPositionEmbeddingsRaw.(float64); castOk { |
| 143 | + model.MaxPositionEmbeddings = int(maxPositionEmbeddings) |
| 144 | + } |
| 145 | + } |
| 146 | + } |
| 147 | + return nil |
20 | 148 | } |
21 | 149 |
|
22 | 150 | func ReshapeOutput(input *[]float32, meta InputOutputInfo, paddingMask [][]bool, sequenceLength int) OutputArray { |
@@ -84,38 +212,3 @@ func flatDataTo3D(input *[]float32, paddingMask [][]bool, sequenceLength int, di |
84 | 212 |
|
85 | 213 | return output |
86 | 214 | } |
87 | | - |
88 | | -func LoadModel(path string, onnxFilename string, options *options.Options) (*Model, error) { |
89 | | - model := &Model{ |
90 | | - Path: path, |
91 | | - OnnxFilename: onnxFilename, |
92 | | - Pipelines: make(map[string]Pipeline), |
93 | | - } |
94 | | - |
95 | | - err := LoadOnnxModelBytes(model) |
96 | | - if err != nil { |
97 | | - return nil, err |
98 | | - } |
99 | | - |
100 | | - err = CreateModelBackend(model, options) |
101 | | - if err != nil { |
102 | | - return nil, err |
103 | | - } |
104 | | - |
105 | | - tkErr := LoadTokenizer(model, options) |
106 | | - if tkErr != nil { |
107 | | - return nil, tkErr |
108 | | - } |
109 | | - |
110 | | - model.Destroy = func() error { |
111 | | - destroyErr := model.Tokenizer.Destroy() |
112 | | - switch options.Backend { |
113 | | - case "ORT": |
114 | | - destroyErr = errors.Join(destroyErr, model.ORTModel.Destroy()) |
115 | | - case "GO", "XLA": |
116 | | - model.GoMLXModel.Destroy() |
117 | | - } |
118 | | - return destroyErr |
119 | | - } |
120 | | - return model, nil |
121 | | -} |
|
0 commit comments