bart

Documentation

Overview

Package bart implements the transformer model introduced by Mike et al., 2019. "BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension" https://arxiv.org/abs/1910.13461

Index

• type Model
  • func New(config config.Config, embeddingsStoragePath string) *Model
  • func (m *Model) Close()
  • func (m *Model) Decode(inputIDs []int, encoderHiddenStates []ag.Node, ...) ([]ag.Node, decoder.KeysValuesPairs)
  • func (m *Model) Encode(inputIDs []int) []ag.Node
  • func (m *Model) Process(inputIDs []int) []ag.Node

Types

type Model

type Model struct {
	nn.BaseModel
	Config     config.Config
	Embeddings *embeddings.Model
	Encoder    *encoder.Model
	Decoder    *decoder.Model
}

Model implements a BART model.

func New

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

New returns a new BART Model.

func (*Model) Close

func (m *Model) Close()

Close closes the BART model's embeddings DB.

func (*Model) Decode

func (m *Model) Decode(
	inputIDs []int,
	encoderHiddenStates []ag.Node,
	pastKeysValuesPairs decoder.KeysValuesPairs,
) ([]ag.Node, decoder.KeysValuesPairs)

Decode performs the BART decoding.

func (*Model) Encode

func (m *Model) Encode(inputIDs []int) []ag.Node

Encode performs the BART encoding.

func (*Model) Process

func (m *Model) Process(inputIDs []int) []ag.Node

Process performs the forward step for each input and returns the result.