gobert: github.com/buckhx/gobert/model Index | Files | Directories

package model

import "github.com/buckhx/gobert/model"

Package model provides functionality for working with exported BERT models


Package Files

bert.go classifier.go embedding.go option.go


const (
    //#	UniqueIDsOp    = "unique_ids"
    InputIDsOp     = "input_ids"
    InputMaskOp    = "input_mask"
    InputTypeIDsOp = "input_type_ids"

Operation names

const (
    DefaultSeqLen    = 128
    DefaultVocabFile = "vocab.txt"

Default values

const (
    ClassifierOutputOp = "probabilities"
    ClassifierModelTag = "bert-tuned"
    ClassifierSeqLen   = 64


const (
    EmbeddingModelTag = "bert-pretrained"
    EmbeddingOp       = "embedding"

Embedding Defaults

func Print Uses

func Print(m *tf.SavedModel)

Print is a utility for printing the operations in a saved model

type Bert Uses

type Bert struct {
    // contains filtered or unexported fields

Bert is a model that translates features to values from an exported model. It processes as follows: Pipeline: text -> FeatureFactory -> TensorFunc -> InputFunc -> ModelFunc -> Value

func NewBert Uses

func NewBert(m *tf.SavedModel, vocabPath string, opts ...BertOption) (Bert, error)

NewBert will create a new default BERT model from the exported model and vocab. Generally used for producing embeddings

func NewBertClassifier Uses

func NewBertClassifier(path string, vocabPath string, opts ...BertOption) (Bert, error)

NewBertClassifier returns a model configured for classification after being fine-tuned with run_classification.py

func NewEmbeddings Uses

func NewEmbeddings(path string, opts ...BertOption) (Bert, error)

NewEmbeddings returns a pre-trained model for text embeddings

func (Bert) Features Uses

func (b Bert) Features(texts ...string) []tokenize.Feature

Features will tokenize a text

func (Bert) PredictValues Uses

func (b Bert) PredictValues(texts ...string) ([]ValueProvider, error)

PredictValues will run the BERT model on the provided texts. The returned values are in the same order as the provided texts.

type BertOption Uses

type BertOption func(b Bert) Bert

BertOption configures a BERT model

func WithFeatureFactory Uses

func WithFeatureFactory(ff *tokenize.FeatureFactory) BertOption

WithFeatureFactory replaces the default feature factory

func WithInputFunc Uses

func WithInputFunc(fn TensorInputFunc) BertOption

WithInputFunc updates the input func, used if input tensors vary from defaults

func WithModelFunc Uses

func WithModelFunc(fn estimator.ModelFunc) BertOption

WithModelFunc applies the given model func, used when outputs do not match the default

func WithSeqLen Uses

func WithSeqLen(l int32) BertOption

WithSeqLen applies the seqlen, should match max_seq_len from trained model

func WithTokenizer Uses

func WithTokenizer(tkz tokenize.VocabTokenizer) BertOption

WithTokenizer applies the given tokenizer to the model

type FeatureTensorFunc Uses

type FeatureTensorFunc func(fs ...tokenize.Feature) (map[string]*tf.Tensor, error)

FeatureTensorFunc translates features to tensors

type TensorInputFunc Uses

type TensorInputFunc func(map[string]*tf.Tensor) estimator.InputFunc

TensorInputFunc maps tensors to an estimator.InputFunc in the Predict pipeline

type ValueProvider Uses

type ValueProvider interface {
    Value() interface{}

ValueProvider is a simple interface for tensors responses without the baggage


estimatorPackage estimator is a utility method for interacting with tf models.

Package model imports 5 packages (graph) and is imported by 4 packages. Updated 2019-07-29. Refresh now. Tools for package owners.