bert

type Config struct {
	Architectures             []string          `json:"architectures"`
	AttentionProbsDropoutProb float64           `json:"attention_probs_dropout_prob"`
	GradientCheckpointing     bool              `json:"gradient_checkpointing"`
	HiddenAct                 string            `json:"hidden_act"`
	HiddenDropoutProb         float64           `json:"hidden_dropout_prob"`
	HiddenSize                int               `json:"hidden_size"`
	EmbeddingsSize            int               `json:"embeddings_size"`
	InitializerRange          float64           `json:"initializer_range"`
	IntermediateSize          int               `json:"intermediate_size"`
	LayerNormEps              float64           `json:"layer_norm_eps"`
	MaxPositionEmbeddings     int               `json:"max_position_embeddings"`
	ModelType                 string            `json:"model_type"`
	NumAttentionHeads         int               `json:"num_attention_heads"`
	NumHiddenLayers           int               `json:"num_hidden_layers"`
	PadTokenId                int               `json:"pad_token_id"`
	PositionEmbeddingType     string            `json:"position_embedding_type"`
	TransformersVersion       string            `json:"transformers_version"`
	TypeVocabSize             int               `json:"type_vocab_size"`
	UseCache                  bool              `json:"use_cache"`
	VocabSize                 int               `json:"vocab_size"`
	ID2Label                  map[string]string `json:"id2label"`
	Cybertron                 struct {
		Training            bool   `json:"training"`
		TokensStoreName     string `json:"tokens_store_name"`
		PositionsStoreName  string `json:"positions_store_name"`
		TokenTypesStoreName string `json:"token_types_store_name"`
	}
}

type ConfigFile interface {
	Config | TokenizerConfig
}

type Embeddings struct {
	nn.Module
	Vocab      *vocabulary.Vocabulary
	Tokens     *emb.Model // string
	Positions  *emb.Model
	TokenTypes *emb.Model
	Norm       *layernorm.Model
	Projector  *linear.Model
	Config     Config
}

type Encoder struct {
	nn.Module
	Layers nn.ModuleList[*EncoderLayer]
	Config Config
}

type EncoderLayer struct {
	nn.Module
	SelfAttention *SelfAttentionBlock
	FF            *FeedForwardBlock
	Config        Config
}

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

type FeedForwardBlockConfig struct {
	Dim        int
	HiddenDim  int
	Activation activation.Activation
}

type Model struct {
	nn.Module
	Embeddings *Embeddings
	Encoder    *Encoder
	Pooler     *Pooler
	Config     Config
}

type ModelForMaskedLM struct {
	nn.Module
	// Bart is the fine-tuned BERT model.
	Bert *Model
	// Layers contains the feedforward layers for masked language modeling.
	Layers nn.ModuleList[nn.StandardModel]
}

type ModelForQuestionAnswering struct {
	nn.Module
	// Bart is the fine-tuned BERT model.
	Bert *Model
	// Classifier is the linear layer for span classification.
	Classifier *linear.Model
}

type ModelForSequenceClassification struct {
	nn.Module
	// Bart is the fine-tuned BERT model.
	Bert *Model
	// Classifier is the linear layer for sequence classification.
	Classifier *linear.Model
}

type ModelForSequenceEncoding struct {
	nn.Module
	// Bart is the fine-tuned BERT model.
	Bert *Model
}

type ModelForTokenClassification struct {
	nn.Module
	// Bart is the fine-tuned BERT model.
	Bert *Model
	// Classifier is the linear layer for sequence classification.
	Classifier *linear.Model
}

type Pooler struct {
	nn.Module
	Model nn.ModuleList[nn.StandardModel]
}

type PoolingStrategyType int

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        int
	NumOfHeads int
}

type TokenizerConfig struct {
	DoLowerCase          bool        `json:"do_lower_case"`
	UnkToken             string      `json:"unk_token"`
	SepToken             string      `json:"sep_token"`
	PadToken             string      `json:"pad_token"`
	ClsToken             string      `json:"cls_token"`
	MaskToken            string      `json:"mask_token"`
	TokenizeChineseChars bool        `json:"tokenize_chinese_chars"`
	StripAccents         interface{} `json:"strip_accents"`
	ModelMaxLength       int         `json:"model_max_length"`
}