Documentation ¶
Overview ¶
Package model provides functionality for working with exported BERT models
Index ¶
Constants ¶
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 )
DefaultOverrides
const ( EmbeddingModelTag = "bert-pretrained" EmbeddingOp = "embedding" )
Embedding Defaults
Variables ¶
This section is empty.
Functions ¶
func Print ¶
func Print(m *tf.SavedModel)
Print is a utility for printing the operations in a saved model
Types ¶
type Bert ¶
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 ¶
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 ¶
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 ¶
func NewEmbeddings(path string, opts ...BertOption) (Bert, error)
NewEmbeddings returns a pre-trained model for text embeddings
func (Bert) PredictValues ¶
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 ¶
BertOption configures a BERT model
func WithFeatureFactory ¶
func WithFeatureFactory(ff *tokenize.FeatureFactory) BertOption
WithFeatureFactory replaces the default feature factory
func WithInputFunc ¶
func WithInputFunc(fn TensorInputFunc) BertOption
WithInputFunc updates the input func, used if input tensors vary from defaults
func WithModelFunc ¶
func WithModelFunc(fn estimator.ModelFunc) BertOption
WithModelFunc applies the given model func, used when outputs do not match the default
func WithSeqLen ¶
func WithSeqLen(l int32) BertOption
WithSeqLen applies the seqlen, should match max_seq_len from trained model
func WithTokenizer ¶
func WithTokenizer(tkz tokenize.VocabTokenizer) BertOption
WithTokenizer applies the given tokenizer to the model
type FeatureTensorFunc ¶
FeatureTensorFunc translates features to tensors
type TensorInputFunc ¶
TensorInputFunc maps tensors to an estimator.InputFunc in the Predict pipeline
type ValueProvider ¶
type ValueProvider interface {
Value() interface{}
}
ValueProvider is a simple interface for tensors responses without the baggage