bart

type Cache [][2]multiheadattention.Cache

type Classifier struct {
	nn.Module
	// Config is the configuration of the classifier.
	Config ClassifierConfig
	// Layers is the list of layers of the MLP.
	Layers nn.ModuleList[nn.StandardModel]
}

type ClassifierConfig struct {
	// InputSize is the input size of the classifier.
	InputSize int
	// HiddenSize is the hidden size of the classifier.
	HiddenSize int
	// OutputSize is the output size of the classifier.
	OutputSize int
	// PoolerDropout is the dropout rate for the classifier.
	PoolerDropout float64
}

type Config struct {
	NumLabels                  int               `json:"_num_labels,omitempty"`
	ActivationDropout          float64           `json:"activation_dropout,omitempty"`
	ActivationFunction         string            `json:"activation_function,omitempty"`
	BiasLogits                 bool              `json:"add_bias_logits,omitempty"`
	FinalLayerNorm             bool              `json:"add_final_layer_norm,omitempty"`
	Architecture               []string          `json:"architectures,omitempty"`
	AttentionDropout           float64           `json:"attention_dropout,omitempty"`
	BosTokenID                 int               `json:"bos_token_id,omitempty"`
	ClassifierDropout          float64           `json:"classif_dropout,omitempty"`
	DModel                     int               `json:"d_model,omitempty"`
	DecoderAttentionHeads      int               `json:"decoder_attention_heads,omitempty"`
	DecoderFFNDim              int               `json:"decoder_ffn_dim,omitempty"`
	DecoderLayerDrop           float64           `json:"decoder_layerdrop,omitempty"`
	DecoderLayers              int               `json:"decoder_layers,omitempty"`
	DecoderStartTokenID        int               `json:"decoder_start_token_id,omitempty"`
	Dropout                    float64           `json:"dropout,omitempty"`
	EncoderAttentionHeads      int               `json:"encoder_attention_heads,omitempty"`
	EncoderFFNDim              int               `json:"encoder_ffn_dim,omitempty"`
	EncoderLayerDrop           float64           `json:"encoder_layerdrop,omitempty"`
	EncoderLayers              int               `json:"encoder_layers,omitempty"`
	EosTokenID                 int               `json:"eos_token_id,omitempty"`
	FineTuningTask             string            `json:"finetuning_task,omitempty"`
	ForceBosTokenToBeGenerated bool              `json:"force_bos_token_to_be_generated,omitempty"`
	ID2Label                   map[string]string `json:"id2label,omitempty"`
	InitStd                    float64           `json:"init_std,omitempty"`
	IsEncoderDecoder           bool              `json:"is_encoder_decoder,omitempty"`
	Label2ID                   map[string]int    `json:"label2id,omitempty"`
	LengthPenalty              float64           `json:"length_penalty,omitempty"`
	MaxPositionEmbeddings      int               `json:"max_position_embeddings,omitempty"`
	ModelType                  string            `json:"model_type,omitempty"`
	NormalizeBefore            bool              `json:"normalize_before,omitempty"`
	NormalizeEmbedding         bool              `json:"normalize_embedding,omitempty"`
	NumHiddenLayers            int               `json:"num_hidden_layers,omitempty"`
	OutputPast                 bool              `json:"output_past,omitempty"`
	PadTokenID                 int               `json:"pad_token_id,omitempty"`
	ScaleEmbedding             bool              `json:"scale_embedding,omitempty"`
	StaticPositionEmbeddings   bool              `json:"static_position_embeddings,omitempty"`
	TotalFlos                  float64           `json:"total_flos,omitempty"`
	VocabSize                  int               `json:"vocab_size,omitempty"`
	NumBeams                   int               `json:"num_beams,omitempty"`
	MaxLength                  int               `json:"max_length,omitempty"`
	MinLength                  int               `json:"min_length,omitempty"`
	BadWordsIDs                [][]int           `json:"bad_words_ids,omitempty"`
	EarlyStopping              bool              `json:"early_stopping,omitempty"`
	NoRepeatNGramSize          int               `json:"no_repeat_ngram_size,omitempty"`
	ExtraSpecialTokens         map[int]string    `json:"extra_special_tokens,omitempty"`
	Cybertron                  struct {
		Training                           bool   `json:"training,omitempty"`
		PositionalEncoderOffset            int    `json:"positional_encoder_offset,omitempty"`
		SharedEmbeddingsStoreName          string `json:"shared_embeddings_store_name,omitempty"`
		DecoderPositionalEncodingStoreName string `json:"decoder_positional_encoding_store_name,omitempty"`
		EncoderPositionalEncodingStoreName string `json:"encoder_positional_encoding_store_name,omitempty"`
	}
}

type CrossAttentionBlock struct {
	nn.Module
	Attention *multiheadattention.Model
	Norm      *layernorm.Model
}

type CrossAttentionBlockConfig struct {
	Dim             int
	NumOfHeads      int
	NormalizeBefore bool
}

type Decoder struct {
	nn.Module
	// Embeddings is the embedding module.
	Embeddings *Embeddings
	// Layers is the list of decoder layers.
	Layers []*DecoderLayer
	// LayerNorm is the layer normalization module.
	LayerNorm *layernorm.Model
	// Config is the configuration of the decoder.
	Config Config
}

type DecoderLayer struct {
	nn.Module
	// SelfAttention is the self-attention block.
	SelfAttention ResidualNormSelfAttention
	// CrossAttention is the cross-attention block.
	CrossAttention ResidualNormCrossAttention
	// FF is the feed-forward block with normalization and residual connection.
	FF ResidualNormFeedForward
	// Config is the configuration of the decoder layer.
	Config Config
}

type DecodingInput struct {
	// InputIDs are the input IDs for the decoder.
	InputIDs []int
	// CurLen is the current length of the generating sequence.
	CurLen int
	// Cache is the cache for the decoder.
	Cache Cache
}

type DecodingOutput struct {
	// LogProbRaw is the raw (not processed) log probability of the generated token.
	LogProbRaw mat.Tensor
	// LogProbValue is the post-processed log probability of the generated token.
	LogProbValue mat.Matrix
	// NextCache is the next cache.
	NextCache Cache
}

type Embeddings struct {
	nn.Module
	// SharedEmbeddings is the shared embedding module.
	SharedEmbeddings embedding.Shared
	// PositionalEncoder is the positional encoder module.
	PositionalEncoder *PositionalEncoder
	// Norm is the normalization module.
	Norm *layernorm.Model
	// ScaleFactor is the scaling factor for the shared embeddings.
	ScaleFactor *nn.Buffer
	// Config is the configuration of the embeddings.
	Config Config
}

type Encoder struct {
	nn.Module
	// Embeddings is the embedding layer.
	Embeddings *Embeddings
	// Layers is the list of encoder layers.
	Layers nn.ModuleList[*EncoderLayer]
	// LayerNorm is the layer normalization module.
	LayerNorm *layernorm.Model
	// Config is the configuration of the encoder.
	Config Config
}

type EncoderLayer struct {
	nn.Module
	// SelfAttention is the self attention block.
	SelfAttention ResidualNormSelfAttention
	// FF is the feed-forward block with normalization and residual connection.
	FF ResidualNormFeedForward
	// Config is the configuration of the encoder layer.
	Config Config
}

type FeedForwardBlock struct {
	nn.Module
	FFN  nn.ModuleList[nn.StandardModel]
	Norm *layernorm.Model
}

type Model struct {
	nn.Module
	// Config is the model configuration.
	Config Config
	// Encoder is the encoder model.
	Encoder *Encoder
	// Decoder is the decoder model.
	Decoder *Decoder
	// Embeddings contains the embeddings shared between the encoder and the decoder.
	Embeddings *embedding.Model
}

type ModelForConditionalGeneration struct {
	nn.Module
	// Bart is the fine-tuned BART model.
	Bart *Model
	// Projection is the projection layer from the decoder output to the vocabulary.
	Projection *linear.Model
	// PadMask is the mask for the pad token.
	PadMask *nn.Buffer
	// EosMask is the mask for the EOS token.
	EosMask *nn.Buffer
}

type ModelForSequenceClassification struct {
	nn.Module
	// Bart is the Bart fine-tuned model.
	Bart *Model
	// Classifier is the final classifier layer.
	Classifier *Classifier
}

type NewFeedForwardBlockConfig struct {
	// Dim is the dimension of the input.
	Dim int
	// HiddenDim is the dimension of the hidden layer.
	HiddenDim int
	// ActivationFunction is the activation function.
	Activation activation.Activation
	// NormalizeBefore is whether to normalize the input before the MLP.
	NormalizeBefore bool
}

type PositionalEncoder struct {
	nn.Module
	// Embeddings contains the embeddings for each position.
	Embeddings *embedding.Model
	// Config contains the configuration settings.
	Config PositionalEncoderConfig
}

type PositionalEncoderConfig struct {
	NumEmbeddings int
	EmbeddingDim  int
	PaddingIDX    int
	Offset        int
	StoreName     string
	Trainable     bool
}

type PostNormCrossAttentionBlock struct {
	*CrossAttentionBlock
}

type PostNormFeedForwardBlock struct {
	*FeedForwardBlock
}

type PostNormSelfAttentionBlock struct {
	*SelfAttentionBlock
}

type PreNormCrossAttentionBlock struct {
	*CrossAttentionBlock
}

type PreNormFeedForwardBlock struct {
	*FeedForwardBlock
}

type PreNormSelfAttentionBlock struct {
	*SelfAttentionBlock
}

type ResidualNormCrossAttention interface {
	// Forward performs the forward pass.
	Forward(cache multiheadattention.Cache, seq1 []mat.Tensor, seq2 []mat.Tensor) ([]mat.Tensor, multiheadattention.Cache)
}

type ResidualNormFeedForward interface {
	Forward(xs []mat.Tensor) []mat.Tensor
}

type ResidualNormSelfAttention interface {
	Forward(cache multiheadattention.Cache, xs []mat.Tensor) ([]mat.Tensor, multiheadattention.Cache)
}

type SelfAttentionBlock struct {
	nn.Module
	// Attention is the multi-head attention module.
	Attention *multiheadattention.Model
	// Norm is the layer normalization module.
	Norm *layernorm.Model
}

type SelfAttentionBlockConfig struct {
	// Dim is the dimension of the input and output.
	Dim int
	// NumOfHeads is the number of heads.
	NumOfHeads int
	// NormalizeBefore indicates whether the normalization is applied before or after the attention.
	NormalizeBefore bool
	// UseCausalMask indicates whether to use a causal mask.
	UseCausalMask bool
}