@@ -82,8 +82,10 @@ func round3decimals(x float64) float64 {
8282
8383func trainSimilarity (t * testing.T ,
8484 config TrainingConfig ,
85- examplesLhs ,
86- examplesRhs []string ) []float64 {
85+ examplesLHS ,
86+ examplesRHS []string ,
87+ ) []float64 {
88+ t .Helper ()
8789 // Create a new GoMLX training session. Currently, training is only possible by loading an onnx model
8890 // into GoMLX, fine-tuning it, and then writing it back to onnx. Hugot deals with the details
8991 // for you here.
@@ -118,7 +120,7 @@ func trainSimilarity(t *testing.T,
118120 }()
119121
120122 // we now load the newly trained onnx model and generate the predictions with onnxruntime backend
121- return runModel (t , "ORT" , examplesLhs , examplesRhs , "./models/testTrain" )
123+ return runModel (t , "ORT" , examplesLHS , examplesRHS , "./models/testTrain" )
122124}
123125
124126func TestTrainSemanticSimilarity (t * testing.T ) {
@@ -130,8 +132,8 @@ func TestTrainSemanticSimilarity(t *testing.T) {
130132 checkT (t , err )
131133 lines := bytes .Split (data , []byte ("\n " ))
132134
133- var examplesLhs []string
134- var examplesRhs []string
135+ var examplesLHS []string
136+ var examplesRHS []string
135137 var scores []float64
136138
137139 for _ , line := range lines {
@@ -140,16 +142,16 @@ func TestTrainSemanticSimilarity(t *testing.T) {
140142 if err != nil {
141143 t .Fatal (err )
142144 }
143- examplesLhs = append (examplesLhs , example ["sentence1" ].(string ))
144- examplesRhs = append (examplesRhs , example ["sentence2" ].(string ))
145+ examplesLHS = append (examplesLHS , example ["sentence1" ].(string ))
146+ examplesRHS = append (examplesRHS , example ["sentence2" ].(string ))
145147 scores = append (scores , example ["score" ].(float64 ))
146148 }
147149
148150 // first we run the untrained onnx model with onnxruntime backend
149- similaritiesOnnxruntime := runModel (t , "ORT" , examplesLhs , examplesRhs , modelPath )
151+ similaritiesOnnxruntime := runModel (t , "ORT" , examplesLHS , examplesRHS , modelPath )
150152
151153 // we do the same for GoMLX and check the forward pass results match
152- similaritiesGoMLX := runModel (t , "XLA" , examplesLhs , examplesRhs , modelPath )
154+ similaritiesGoMLX := runModel (t , "XLA" , examplesLHS , examplesRHS , modelPath )
153155
154156 for i := range similaritiesOnnxruntime {
155157 assert .Equal (t , round3decimals (similaritiesOnnxruntime [i ]), round3decimals (similaritiesGoMLX [i ]))
@@ -185,7 +187,7 @@ func TestTrainSemanticSimilarity(t *testing.T) {
185187 },
186188 Verbose : true ,
187189 }
188- similaritiesGoMLXTrained := trainSimilarity (t , trainingConfig , examplesLhs , examplesRhs )
190+ similaritiesGoMLXTrained := trainSimilarity (t , trainingConfig , examplesLHS , examplesRHS )
189191
190192 fmt .Println ("GoMLX trained model predictions:" )
191193 for i := range similaritiesGoMLXTrained {
@@ -202,24 +204,24 @@ func TestTrainSemanticSimilarity(t *testing.T) {
202204
203205 // we can also train a model using an in memory dataset. For this we create the slice of examples manually.
204206 var examples []datasets.SemanticSimilarityExample
205- for i := 0 ; i < len (examplesLhs ); i ++ {
207+ for i := 0 ; i < len (examplesLHS ); i ++ {
206208 examples = append (examples , datasets.SemanticSimilarityExample {
207- Sentence1 : examplesLhs [i ],
208- Sentence2 : examplesRhs [i ],
209+ Sentence1 : examplesLHS [i ],
210+ Sentence2 : examplesRHS [i ],
209211 Score : float32 (scores [i ]),
210212 })
211213 }
212214 inMemoryDataset , err := datasets .NewInMemorySemanticSimilarityDataset (examples , 1 , nil )
213215 checkT (t , err )
214216 trainingConfig .TrainDataset = inMemoryDataset
215- similaritiesGoMLXTrainedInMemory := trainSimilarity (t , trainingConfig , examplesLhs , examplesRhs )
217+ similaritiesGoMLXTrainedInMemory := trainSimilarity (t , trainingConfig , examplesLHS , examplesRHS )
216218 for i := range similaritiesGoMLXTrainedInMemory {
217219 assert .Equal (t , round3decimals (similaritiesGoMLXTrained [i ]), round3decimals (similaritiesGoMLXTrainedInMemory [i ]))
218220 }
219221
220222 // we can also freeze layers
221223 trainingConfig .Options = append (trainingConfig .Options , WithFreezeLayers ([]int {- 1 })) // freeze all layers but the last one
222- similaritiesGoMLXTrainedFrozen := trainSimilarity (t , trainingConfig , examplesLhs , examplesRhs )
224+ similaritiesGoMLXTrainedFrozen := trainSimilarity (t , trainingConfig , examplesLHS , examplesRHS )
223225
224226 fmt .Println ("GoMLX trained model predictions freezing all layers but the last one:" )
225227 for i := range similaritiesGoMLXTrainedFrozen {
@@ -230,7 +232,7 @@ func TestTrainSemanticSimilarity(t *testing.T) {
230232func rmse (predictions []float64 , labels []float64 ) float64 {
231233 var sum float64
232234 for i := 0 ; i < len (predictions ); i ++ {
233- sum += math . Pow (predictions [i ]- labels [i ], 2 )
235+ sum += (predictions [i ] - labels [i ]) * ( predictions [ i ] - labels [ i ] )
234236 }
235237 return math .Sqrt (sum / float64 (len (predictions )))
236238}
0 commit comments