barthead

Documentation

Index

• type Classification
  • func NewClassification(config ClassificationConfig) *Classification
• type ClassificationConfig
• type SequenceClassification
  • func LoadModelForSequenceClassification(modelPath string) (*SequenceClassification, error)
  • func NewSequenceClassification(config bartconfig.Config, embeddingsPath string) *SequenceClassification
  • func (m *SequenceClassification) Classify(in interface{}) ag.Node
  • func (m *SequenceClassification) Close()

Types

type Classification

type Classification struct {
	Config ClassificationConfig
	*stack.Model
}

Classification is a model for BART head for sentence-level classification tasks.

func NewClassification

func NewClassification(config ClassificationConfig) *Classification

NewClassification returns a new Classification.

type ClassificationConfig

type ClassificationConfig struct {
	InputSize     int
	HiddenSize    int
	OutputSize    int
	PoolerDropout float64
}

ClassificationConfig provides configuration settings for a BART head for sentence-level Classification model.

type SequenceClassification

type SequenceClassification struct {
	nn.BaseModel
	BART           *bart.Model
	Classification *Classification
}

SequenceClassification is a model for sentence-level classification tasks which embeds a BART pre-trained model.

func LoadModelForSequenceClassification

func LoadModelForSequenceClassification(modelPath string) (*SequenceClassification, error)

LoadModelForSequenceClassification loads a SequenceClassification model from file.

func NewSequenceClassification

func NewSequenceClassification(config bartconfig.Config, embeddingsPath string) *SequenceClassification

NewSequenceClassification returns a new SequenceClassification.

func (*SequenceClassification) Classify

func (m *SequenceClassification) Classify(in interface{}) ag.Node

Classify performs the classification using the last transformed state.

func (*SequenceClassification) Close

func (m *SequenceClassification) Close()

Close closes the BART model's embeddings DB.