Fix compilation issues with XLA disabled

Author: RJKeevil

datasets/dataset_xla_disabled.go

@@ -4,10 +4,14 @@ package datasets
 
 import "github.com/knights-analytics/hugot/pipelineBackends"
 
-type Dataset interface{}
+type Dataset interface {
+	SetVerbose(bool)
+}
 
 type SemanticSimilarityDataset struct{}
 
 func (s *SemanticSimilarityDataset) SetTokenizationPipeline(_ pipelineBackends.Pipeline) error {
 	return nil
 }
+
+func (s *SemanticSimilarityDataset) SetVerbose(_ bool) {}

hugot_training.go

@@ -97,18 +97,17 @@ func (s *TrainingSession) Train() error {
 func (s *TrainingSession) Save(path string) error {
 	model := s.pipeline.GetModel()
 	if model != nil {
-		xlaModel := model.XLAModel
-		if xlaModel != nil {
-			if err := xlaModel.OnnxModel.ContextToONNX(xlaModel.Ctx); err != nil {
-				return err
+		if s.runtime == "XLA" {
+			xlaModel := model.XLAModel
+			if xlaModel != nil {
+				return xlaModel.Save(path)
+			} else {
+				return fmt.Errorf("xla model is nil")
 			}
-			if err := xlaModel.OnnxModel.SaveToFile(path); err != nil {
-				return err
-			}
-			return nil
 		} else {
-			return fmt.Errorf("xla model is nil")
+			return fmt.Errorf("XLA runtime is required for saving a training model")
 		}
+
 	} else {
 		return fmt.Errorf("pipeline model is nil")
 	}

pipelineBackends/model_xla.go

@@ -227,3 +227,13 @@ func nextPowerOf2(n int) int {
 	}
 	return pow
 }
+
+func (xlaModel *XLAModel) Save(path string) error {
+	if err := xlaModel.OnnxModel.ContextToONNX(xlaModel.Ctx); err != nil {
+		return err
+	}
+	if err := xlaModel.OnnxModel.SaveToFile(path); err != nil {
+		return err
+	}
+	return nil
+}

pipelineBackends/model_xla_disabled.go

@@ -14,10 +14,14 @@ func createXLAModelBackend(_ *Model, _ *options.Options) error {
 	return nil
 }
 
-func createInputTensorsXLA(_ *PipelineBatch, _ []InputOutputInfo) error {
+func createInputTensorsXLA(_ *PipelineBatch, _ []InputOutputInfo, _ bool) error {
 	return nil
 }
 
 func runXLASessionOnBatch(_ *PipelineBatch, _ *BasePipeline) error {
 	return nil
 }
+
+func (xlaModel *XLAModel) Save(_ string) error {
+	return nil
+}