bart

Documentation

Index

• type BPETokenizer
  • func (m *BPETokenizer) Detokenize(tokenIds []int, stripPaddingTokens bool) string
  • func (m *BPETokenizer) Tokenize(text string) ([]int, error)
• type SentencePieceTokenizer
  • func (m *SentencePieceTokenizer) Detokenize(tokenIds []int, stripPaddingTokens bool) string
  • func (m *SentencePieceTokenizer) Tokenize(text string) ([]int, error)
• type Text2Text
  • func LoadText2Text(modelPath string) (*Text2Text, error)
  • func (m *Text2Text) Close() error
  • func (m *Text2Text) Generate(ctx context.Context, text string, opts *text2text.Options) (text2text.Response, error)
• type Tokenizer

Types

type BPETokenizer

type BPETokenizer struct {
	*bpetokenizer.BPETokenizer
	EosTokenID           int
	BosTokenID           int
	PadTokenID           int
	DecoderStartTokenID  int
	ExtraSpecialTokenIDs map[int]string
}

func (*BPETokenizer) Detokenize

func (m *BPETokenizer) Detokenize(tokenIds []int, stripPaddingTokens bool) string

Detokenize returns the text of the input token IDs removing the padding token.

func (*BPETokenizer) Tokenize

func (m *BPETokenizer) Tokenize(text string) ([]int, error)

Tokenize returns the token IDs of the input text applying the EOS pad token.

type SentencePieceTokenizer

type SentencePieceTokenizer struct {
	*sentencepiece.Tokenizer
	EosTokenID          int
	BosTokenID          int
	PadTokenID          int
	DecoderStartTokenID int
}

func (*SentencePieceTokenizer) Detokenize

func (m *SentencePieceTokenizer) Detokenize(tokenIds []int, stripPaddingTokens bool) string

Detokenize returns the text of the input token IDs removing the padding token.

func (*SentencePieceTokenizer) Tokenize

func (m *SentencePieceTokenizer) Tokenize(text string) ([]int, error)

Tokenize returns the token IDs of the input text applying the EOS pad token.

type Text2Text

type Text2Text struct {
	// Model is the model used for conditional generation.
	Model *bart.ModelForConditionalGeneration
	// Tokenizer is the tokenizer used for conditional generation.
	Tokenizer Tokenizer
	// contains filtered or unexported fields
}

Text2Text contains the ModelForConditionalGeneration and the Tokenizer used for conditional generation tasks. For example, Machine Translation and Summarization.

func LoadText2Text

func LoadText2Text(modelPath string) (*Text2Text, error)

LoadText2Text returns a Text2Text loading the model, the embeddings and the tokenizer from a directory.

func (*Text2Text) Close

func (m *Text2Text) Close() error

Close finalizes the Text2Text resources. It satisfies the interface io.Closer.

func (*Text2Text) Generate

func (m *Text2Text) Generate(ctx context.Context, text string, opts *text2text.Options) (text2text.Response, error)

Generate generates a text from the input.

type Tokenizer

type Tokenizer interface {
	Tokenize(text string) ([]int, error)
	Detokenize(tokenIds []int, stripPaddingTokens bool) string
}