sequenceclassification

Documentation

Index

• type Classifier
  • func NewClassifier(config ClassifierConfig) *Classifier
• type ClassifierConfig
• type Model
  • func New(config config.Config, embeddingsPath string) *Model
  • func (m *Model) Classify(inputIds []int) ag.Node
  • func (m *Model) Close()

Types

type Classifier

type Classifier struct {
	Config ClassifierConfig
	*stack.Model
}

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

func NewClassifier

func NewClassifier(config ClassifierConfig) *Classifier

NewClassifier returns a new Classifier.

type ClassifierConfig

type ClassifierConfig struct {
	InputSize     int
	HiddenSize    int
	OutputSize    int
	PoolerDropout mat.Float
}

ClassifierConfig provides configuration settings for a BART head for sentence-level Classifier model.

type Model

type Model struct {
	nn.BaseModel
	BART       *bart.Model
	Classifier *Classifier
}

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

func New

func New(config config.Config, embeddingsPath string) *Model

New returns a new Model for sentence-level classification.

func (*Model) Classify

func (m *Model) Classify(inputIds []int) ag.Node

Classify performs the classification using the last transformed state.

func (*Model) Close

func (m *Model) Close()

Close closes the BART model's embeddings DB.