backends

type BasePipeline struct {
	Model           *Model
	PipelineTimings *timings
	PipelineName    string
	Runtime         string
}

type GoMLXModel struct {
	Backend         backends.Backend
	OnnxModel       onnx.Model
	Ctx             *context.Context // ctx with the model's weights.
	Exec            *context.Exec    // exec is used to execute the model with a context.
	Call            func(ctx *context.Context, inputs []*graph.Node) []*graph.Node
	Destroy         func()
	BatchBuckets    []int // BatchBuckets defines bucket sizes for batch dimension padding.
	SequenceBuckets []int // SequenceBuckets defines bucket sizes for sequence length padding.
	MaxCache        int   // MaxCache sets the maximum number of unique input shapes to cache.
}

type GoTokenizer struct {
	Tokenizer *tokenizer.Tokenizer
}

type InputOutputInfo struct {
	// The name of the input or output
	Name string
	// The input or output's dimensions, if it's a tensor. This should be
	// ignored for non-tensor types.
	Dimensions Shape
}

type Message struct {
	Role      string   `json:"role"`
	Content   string   `json:"content"`
	ImageURLs []string `json:"image_urls,omitempty"` // File paths or data URIs for multimodal support
}

type Model struct {
	ID                    string
	ORTModel              *ORTModel
	GoMLXModel            *GoMLXModel
	Tokenizer             *Tokenizer
	Destroy               func() error
	Pipelines             map[string]Pipeline
	IDLabelMap            map[int]string
	SeparatorToken        string
	Path                  string
	OnnxFilename          string
	OnnxPath              string
	OnnxReader            io.ReadCloser
	InputsMeta            []InputOutputInfo
	OutputsMeta           []InputOutputInfo
	MaxPositionEmbeddings int
	IsGenerative          bool
}

type ORTModel struct {
	Destroy           func() error
	GenerativeSession disabledGenerativeSession // placeholder when ORT disabled
}

type OutputInfo struct {
	Name       string
	Dimensions []int64
}

type Pipeline interface {
	GetStatistics() PipelineStatistics         // Get the pipeline running statistics
	Validate() error                           // Validate the pipeline for correctness
	GetMetadata() PipelineMetadata             // Return metadata information for the pipeline
	GetModel() *Model                          // Return the model used by the pipeline
	IsGenerative() bool                        // Return whether the pipeline is generative
	Run([]string) (PipelineBatchOutput, error) // Run the pipeline on an input
}

type PipelineBatch struct {
	InputValues       any
	DestroyInputs     func() error
	Input             []TokenizedInput
	PaddingMask       [][]bool
	OutputValues      []any
	Size              int
	MaxSequenceLength int
	MaxNewTokens      int
	// Multimodal support
	Images            any // Will hold *ortgenai.Images for generative models
	DestroyMultimodal func() error
}

type PipelineBatchOutput interface {
	GetOutput() []any
}

type PipelineConfig[T Pipeline] struct {
	ModelPath    string
	Name         string
	OnnxFilename string
	Options      []PipelineOption[T]
}

type PipelineMetadata struct {
	OutputsInfo []OutputInfo
}

type PipelineOption[T Pipeline] func(eo T) error

type PipelineStatistics struct {
	TokenizerTotalTime             time.Duration
	TokenizerExecutionCount        uint64
	TokenizerAvgQueryTime          time.Duration
	OnnxTotalTime                  time.Duration
	OnnxExecutionCount             uint64
	OnnxAvgQueryTime               time.Duration
	TotalQueries                   uint64
	TotalDocuments                 uint64
	AverageLatency                 time.Duration
	AverageBatchSize               float64
	FilteredResults                uint64
	AvgPrefillSeconds              float64
	TokensPerSecond                float64
	CumulativePrefillSum           float64
	CumulativePrefillCount         int
	CumulativeTokens               int
	CumulativeTokenDurationSeconds float64
}

type RustTokenizer struct{}

type SequenceDelta struct {
	Token string
	Index int
}

type Shape []int64

type TokenizedInput struct {
	Raw               string
	Tokens            []string
	TokenIDs          []uint32
	TypeIDs           []uint32
	AttentionMask     []uint32
	SpecialTokensMask []uint32
	Offsets           [][2]uint
	WordIDs           []int
	MaxAttentionIndex int
}

type Tokenizer struct {
	RustTokenizer    *RustTokenizer
	GoTokenizer      *GoTokenizer
	TokenizerTimings *timings
	Destroy          func() error
	Runtime          string
	MaxAllowedTokens int
}