Skip to content

Commit 4e57711

Browse files
RJKeevilRJKeevil
committed
Better handling for disabled backends
1 parent a21d7d6 commit 4e57711

14 files changed

Lines changed: 241 additions & 194 deletions

‎Dockerfile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#--- dockerfile to test hugot ---
22

3-
ARG GO_VERSION=1.24.0
3+
ARG GO_VERSION=1.24.1
44
ARG ONNXRUNTIME_VERSION=1.20.1
55
ARG BUILD_PLATFORM=linux/amd64
66

‎cmd/main.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
//go:build !NOORT || ALL
2+
13
package main
24

35
import (

‎cmd/main_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
//go:build !NOORT || ALL
2+
13
package main
24

35
import (

‎cuda.Dockerfile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#--- dockerfile to test hugot ---
22

3-
ARG GO_VERSION=1.24.0
3+
ARG GO_VERSION=1.24.1
44
ARG ONNXRUNTIME_VERSION=1.20.1
55
ARG BUILD_PLATFORM=linux/amd64
66

‎datasets/dataset.go‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
package datasets
2+
3+
import (
4+
"bufio"
5+
"encoding/json"
6+
"fmt"
7+
"io"
8+
"path/filepath"
9+
10+
"github.com/knights-analytics/hugot/pipelineBackends"
11+
"github.com/knights-analytics/hugot/pipelines"
12+
"github.com/knights-analytics/hugot/util"
13+
)
14+
15+
func (s *SemanticSimilarityDataset) SetVerbose(v bool) {
16+
s.verbose = v
17+
}
18+
19+
func (s *SemanticSimilarityDataset) SetTokenizationPipeline(pipeline pipelineBackends.Pipeline) error {
20+
if pipeline == nil {
21+
return fmt.Errorf("pipeline is required")
22+
}
23+
if _, ok := pipeline.(*pipelines.FeatureExtractionPipeline); !ok {
24+
return fmt.Errorf("tokenization pipeline must be a FeatureExtractionPipeline")
25+
}
26+
s.pipeline = pipeline.(*pipelines.FeatureExtractionPipeline)
27+
return nil
28+
}
29+
30+
func (s *SemanticSimilarityDataset) Validate() error {
31+
if len(s.trainingExamples) == 0 {
32+
if s.trainingPath == "" {
33+
return fmt.Errorf("training path is required")
34+
}
35+
if filepath.Ext(s.trainingPath) != ".jsonl" {
36+
return fmt.Errorf("training path must be a .jsonl file")
37+
}
38+
}
39+
return nil
40+
}
41+
42+
// SemanticSimilarityExample is a single example for the semantic similarity dataset.
43+
type SemanticSimilarityExample struct {
44+
Sentence1 string `json:"sentence1"`
45+
Sentence2 string `json:"sentence2"`
46+
Score float32 `json:"score"`
47+
Data map[string]any // to store any additional data for the example. Not used by the dataset.
48+
}
49+
50+
type ExamplePreprocessFunc func([]SemanticSimilarityExample) ([]SemanticSimilarityExample, error)
51+
52+
// NewSemanticSimilarityDataset creates a new SemanticSimilarityDataset.
53+
// The trainingPath must be a .jsonl file where each line has the following format:
54+
// {"sentence1":"A plane is taking off.","sentence2":"An air plane is taking off.","score":1.0}
55+
// The score is a float value between 0 and 1.
56+
// preprocessFunc here must be a function that takes a slice of SemanticSimilarityExample and returns a slice of SemanticSimilarityExample.
57+
// This function can be used to apply any custom preprocessing to the example batch before they are passed to the model.
58+
func NewSemanticSimilarityDataset(trainingPath string, batchSize int, preprocessFunc ExamplePreprocessFunc) (*SemanticSimilarityDataset, error) {
59+
d := &SemanticSimilarityDataset{
60+
trainingPath: trainingPath,
61+
batchSize: batchSize,
62+
preprocessFunc: preprocessFunc,
63+
}
64+
if err := d.Validate(); err != nil {
65+
return nil, err
66+
}
67+
68+
sourceReadCloser, err := util.OpenFile(trainingPath)
69+
if err != nil {
70+
return nil, err
71+
}
72+
d.reader = bufio.NewReader(sourceReadCloser)
73+
d.sourceFile = sourceReadCloser
74+
return d, nil
75+
}
76+
77+
// NewInMemorySemanticSimilarityDataset creates a new SemanticSimilarityDataset in memory from a slice of examples.
78+
// preprocessFunc here must be a function that takes a slice of SemanticSimilarityExample and returns a slice of SemanticSimilarityExample.
79+
// This function can be used to apply any custom preprocessing to the example batch before they are passed to the model.
80+
func NewInMemorySemanticSimilarityDataset(examples []SemanticSimilarityExample, batchSize int, preprocessFunc ExamplePreprocessFunc) (*SemanticSimilarityDataset, error) {
81+
d := &SemanticSimilarityDataset{
82+
trainingExamples: examples,
83+
batchSize: batchSize,
84+
preprocessFunc: preprocessFunc,
85+
}
86+
if err := d.Validate(); err != nil {
87+
return nil, err
88+
}
89+
return d, nil
90+
}
91+
92+
// Reset resets the dataset to the beginning of the training data (after the epoch is done).
93+
func (s *SemanticSimilarityDataset) Reset() {
94+
if s.verbose {
95+
fmt.Printf("completed epoch in %d batches of %d examples, resetting dataset\n", s.batchN, s.batchSize)
96+
}
97+
s.batchN = 0
98+
99+
if len(s.trainingExamples) == 0 {
100+
if err := s.sourceFile.Close(); err != nil {
101+
panic(err)
102+
}
103+
104+
sourceReadCloser, err := util.OpenFile(s.trainingPath)
105+
if err != nil {
106+
panic(err) // note: these panics will be catched later with the TryExcept
107+
}
108+
s.sourceFile = sourceReadCloser
109+
110+
// restart the reader
111+
s.reader = bufio.NewReader(sourceReadCloser)
112+
}
113+
}
114+
115+
// YieldRaw returns the next raw batch of examples from the dataset. Note that if a preprocessing function has been
116+
// provided at creation time, the examples will be preprocessed before being returned.
117+
func (s *SemanticSimilarityDataset) YieldRaw() ([]SemanticSimilarityExample, error) {
118+
batchCounter := 0
119+
120+
var lineBytes []byte
121+
var readErr error
122+
var lineData SemanticSimilarityExample
123+
124+
examplesBatch := make([]SemanticSimilarityExample, 0, s.batchSize)
125+
var preprocessErr error
126+
127+
for batchCounter < s.batchSize {
128+
if len(s.trainingExamples) > 0 {
129+
// in memory dataset
130+
start := s.batchN * s.batchSize
131+
if start >= len(s.trainingExamples) {
132+
return examplesBatch, io.EOF // return error for reset
133+
}
134+
end := start + s.batchSize
135+
for i := start; i < end && i < len(s.trainingExamples); i++ {
136+
examplesBatch = append(examplesBatch, s.trainingExamples[i])
137+
}
138+
} else {
139+
lineBytes, readErr = util.ReadLine(s.reader)
140+
if readErr != nil {
141+
return examplesBatch, readErr
142+
}
143+
144+
if err := json.Unmarshal(lineBytes, &lineData); err != nil {
145+
return nil, fmt.Errorf("failed to parse JSON line: %w", err)
146+
}
147+
examplesBatch = append(examplesBatch, lineData)
148+
}
149+
batchCounter++
150+
}
151+
s.batchN++
152+
153+
if s.preprocessFunc != nil {
154+
examplesBatch, preprocessErr = s.preprocessFunc(examplesBatch)
155+
if preprocessErr != nil {
156+
return nil, preprocessErr
157+
}
158+
}
159+
160+
return examplesBatch, nil
161+
}
162+
163+
func (s *SemanticSimilarityDataset) Close() error {
164+
if s.sourceFile != nil {
165+
return s.sourceFile.Close()
166+
}
167+
return nil
168+
}

‎datasets/dataset_xla.go‎

Lines changed: 0 additions & 158 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,16 @@ package datasets
44

55
import (
66
"bufio"
7-
"encoding/json"
87
"errors"
98
"fmt"
109
"io"
11-
"path/filepath"
1210
"slices"
1311

1412
"github.com/gomlx/gomlx/ml/train"
1513
"github.com/gomlx/gomlx/types/tensors"
1614

1715
"github.com/knights-analytics/hugot/pipelineBackends"
1816
"github.com/knights-analytics/hugot/pipelines"
19-
"github.com/knights-analytics/hugot/util"
2017
)
2118

2219
type Dataset interface {
@@ -41,154 +38,6 @@ type SemanticSimilarityDataset struct {
4138
verbose bool
4239
}
4340

44-
func (s *SemanticSimilarityDataset) SetVerbose(v bool) {
45-
s.verbose = v
46-
}
47-
48-
func (s *SemanticSimilarityDataset) SetTokenizationPipeline(pipeline pipelineBackends.Pipeline) error {
49-
if pipeline == nil {
50-
return fmt.Errorf("pipeline is required")
51-
}
52-
if _, ok := pipeline.(*pipelines.FeatureExtractionPipeline); !ok {
53-
return fmt.Errorf("tokenization pipeline must be a FeatureExtractionPipeline")
54-
}
55-
s.pipeline = pipeline.(*pipelines.FeatureExtractionPipeline)
56-
return nil
57-
}
58-
59-
func (s *SemanticSimilarityDataset) Validate() error {
60-
if len(s.trainingExamples) == 0 {
61-
if s.trainingPath == "" {
62-
return fmt.Errorf("training path is required")
63-
}
64-
if filepath.Ext(s.trainingPath) != ".jsonl" {
65-
return fmt.Errorf("training path must be a .jsonl file")
66-
}
67-
}
68-
return nil
69-
}
70-
71-
// SemanticSimilarityExample is a single example for the semantic similarity dataset.
72-
type SemanticSimilarityExample struct {
73-
Sentence1 string `json:"sentence1"`
74-
Sentence2 string `json:"sentence2"`
75-
Score float32 `json:"score"`
76-
Data map[string]any // to store any additional data for the example. Not used by the dataset.
77-
}
78-
79-
type ExamplePreprocessFunc func([]SemanticSimilarityExample) ([]SemanticSimilarityExample, error)
80-
81-
// NewSemanticSimilarityDataset creates a new SemanticSimilarityDataset.
82-
// The trainingPath must be a .jsonl file where each line has the following format:
83-
// {"sentence1":"A plane is taking off.","sentence2":"An air plane is taking off.","score":1.0}
84-
// The score is a float value between 0 and 1.
85-
// preprocessFunc here must be a function that takes a slice of SemanticSimilarityExample and returns a slice of SemanticSimilarityExample.
86-
// This function can be used to apply any custom preprocessing to the example batch before they are passed to the model.
87-
func NewSemanticSimilarityDataset(trainingPath string, batchSize int, preprocessFunc ExamplePreprocessFunc) (*SemanticSimilarityDataset, error) {
88-
d := &SemanticSimilarityDataset{
89-
trainingPath: trainingPath,
90-
batchSize: batchSize,
91-
preprocessFunc: preprocessFunc,
92-
}
93-
if err := d.Validate(); err != nil {
94-
return nil, err
95-
}
96-
97-
sourceReadCloser, err := util.OpenFile(trainingPath)
98-
if err != nil {
99-
return nil, err
100-
}
101-
d.reader = bufio.NewReader(sourceReadCloser)
102-
d.sourceFile = sourceReadCloser
103-
return d, nil
104-
}
105-
106-
// NewInMemorySemanticSimilarityDataset creates a new SemanticSimilarityDataset in memory from a slice of examples.
107-
// preprocessFunc here must be a function that takes a slice of SemanticSimilarityExample and returns a slice of SemanticSimilarityExample.
108-
// This function can be used to apply any custom preprocessing to the example batch before they are passed to the model.
109-
func NewInMemorySemanticSimilarityDataset(examples []SemanticSimilarityExample, batchSize int, preprocessFunc ExamplePreprocessFunc) (*SemanticSimilarityDataset, error) {
110-
d := &SemanticSimilarityDataset{
111-
trainingExamples: examples,
112-
batchSize: batchSize,
113-
preprocessFunc: preprocessFunc,
114-
}
115-
if err := d.Validate(); err != nil {
116-
return nil, err
117-
}
118-
return d, nil
119-
}
120-
121-
// Reset resets the dataset to the beginning of the training data (after the epoch is done).
122-
func (s *SemanticSimilarityDataset) Reset() {
123-
if s.verbose {
124-
fmt.Printf("completed epoch in %d batches of %d examples, resetting dataset\n", s.batchN, s.batchSize)
125-
}
126-
s.batchN = 0
127-
128-
if len(s.trainingExamples) == 0 {
129-
if err := s.sourceFile.Close(); err != nil {
130-
panic(err)
131-
}
132-
133-
sourceReadCloser, err := util.OpenFile(s.trainingPath)
134-
if err != nil {
135-
panic(err) // note: these panics will be catched later with the TryExcept
136-
}
137-
s.sourceFile = sourceReadCloser
138-
139-
// restart the reader
140-
s.reader = bufio.NewReader(sourceReadCloser)
141-
}
142-
}
143-
144-
// YieldRaw returns the next raw batch of examples from the dataset. Note that if a preprocessing function has been
145-
// provided at creation time, the examples will be preprocessed before being returned.
146-
func (s *SemanticSimilarityDataset) YieldRaw() ([]SemanticSimilarityExample, error) {
147-
batchCounter := 0
148-
149-
var lineBytes []byte
150-
var readErr error
151-
var lineData SemanticSimilarityExample
152-
153-
examplesBatch := make([]SemanticSimilarityExample, 0, s.batchSize)
154-
var preprocessErr error
155-
156-
for batchCounter < s.batchSize {
157-
if len(s.trainingExamples) > 0 {
158-
// in memory dataset
159-
start := s.batchN * s.batchSize
160-
if start >= len(s.trainingExamples) {
161-
return examplesBatch, io.EOF // return error for reset
162-
}
163-
end := start + s.batchSize
164-
for i := start; i < end && i < len(s.trainingExamples); i++ {
165-
examplesBatch = append(examplesBatch, s.trainingExamples[i])
166-
}
167-
} else {
168-
lineBytes, readErr = util.ReadLine(s.reader)
169-
if readErr != nil {
170-
return examplesBatch, readErr
171-
}
172-
173-
if err := json.Unmarshal(lineBytes, &lineData); err != nil {
174-
return nil, fmt.Errorf("failed to parse JSON line: %w", err)
175-
}
176-
examplesBatch = append(examplesBatch, lineData)
177-
}
178-
batchCounter++
179-
}
180-
s.batchN++
181-
182-
if s.preprocessFunc != nil {
183-
examplesBatch, preprocessErr = s.preprocessFunc(examplesBatch)
184-
if preprocessErr != nil {
185-
return nil, preprocessErr
186-
}
187-
}
188-
189-
return examplesBatch, nil
190-
}
191-
19241
// Yield returns the next batch of examples from the dataset. The examples are tokenized and converted to tensors for the training process.
19342
func (s *SemanticSimilarityDataset) Yield() (spec any, inputs []*tensors.Tensor, labels []*tensors.Tensor, err error) {
19443
exampleBatch, rawErr := s.YieldRaw()
@@ -233,10 +82,3 @@ func (s *SemanticSimilarityDataset) Yield() (spec any, inputs []*tensors.Tensor,
23382

23483
return nil, slices.Concat(inputLhs, inputRhs), labelTensor, rawErr
23584
}
236-
237-
func (s *SemanticSimilarityDataset) Close() error {
238-
if s.sourceFile != nil {
239-
return s.sourceFile.Close()
240-
}
241-
return nil
242-
}

0 commit comments

Comments
 (0)