Documentation
¶
Overview ¶
Package bucket implements a streaming tokenizer of sentences into buckets (or batches) of discrete sizes, to minimize padding.
Fundamentally, it takes as input a channel of sentences and their references (of type `any`, opaque to this library), and it sends on an output channel the sentences tokenized/padded to the same length in the corresponding buckets, along with their corresponding buckets. The order is lost because of the buffering of buckets until they fill.
Example usage:
b := bucket.New(tokenizer).
ByPower(32, 8, 2).
WithMaxDelay(100*time.Millisecond, true)
input := make(chan bucket.SentenceRef)
output := make(chan bucket.Bucket, 10)
go func() {
defer close(input)
for i, text := range sentences {
input <- bucket.SentenceRef{Sentence: text, Reference: i}
}
}()
go b.Run(input, output)
for bk := range output {
// Process bk.Batch, bk.Shape, bk.References...
}
Index ¶
- func TwoBitBucketLen(unpaddedLen int) int
- type Bucket
- type Bucketizer
- func (b *Bucketizer) ByPower(batchSize, minSentenceLength int, base float64) *Bucketizer
- func (b *Bucketizer) ByPowerBudget(tokensBudget, minSentenceLength int, base float64) *Bucketizer
- func (b *Bucketizer) ByTwoBitBucket(batchSize, minSentenceLength int) *Bucketizer
- func (b *Bucketizer) ByTwoBitBucketBudget(tokensBudget, minSentenceLength int) *Bucketizer
- func (b *Bucketizer) Run(input <-chan SentenceRef, output chan<- Bucket)
- func (b *Bucketizer) WithBatchPadding(useBatchPadding bool) *Bucketizer
- func (b *Bucketizer) WithMaxDelay(d time.Duration, useBatchPadding bool) *Bucketizer
- func (b *Bucketizer) WithMaxParallelization(maxParallelization int) *Bucketizer
- func (b *Bucketizer) WithShapeFn(shapeFn ShapeFn) *Bucketizer
- type Reference
- type SentenceRef
- type Shape
- type ShapeFn
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
func TwoBitBucketLen ¶
TwoBitBucketLen returns the smallest size >= unpaddedLen that uses only the two highest bits (either 2^n or 1.5 * 2^n).
It is used by Bucketizer.ByTwoBitBucket and Bucketizer.ByTwoBitBucketBudget.
Types ¶
type Bucket ¶
type Bucket struct {
// Shape of this bucket.
Shape
// Batch, with [Shape.BatchSize x Shape.SentenceLength] tokens, padded.
Batch []int
// References, with Shape.BatchSize references.
References []Reference
// NonPadTokens is the number of non-pad tokens in the batch.
NonPadTokens int
// Error, if any, that occurred during tokenization.
// Errors are returned in individual buckets, with Shape.BatchSize == 1,
// and References set to the failing SentenceRef.
Error error
}
Bucket contains the tokenized sentences and their references.
type Bucketizer ¶
type Bucketizer struct {
// contains filtered or unexported fields
}
Bucketizer represents the streaming bucketizer.
func New ¶
func New(tokenizer api.Tokenizer) *Bucketizer
New creates a new Bucketizer.
By default it starts with a shape function grouping up to 32 elements with lengths rounded up to the lowest power of 2, starting with a minimum of 8 tokens (so 8, 16, 32, ...). It also sets the maximum tokenization parallelization to the number of available CPUs.
Example usage:
b := bucket.New(tokenizer).
ByPower(32, 8, 2).
WithMaxDelay(100*time.Millisecond, true)
input := make(chan bucket.SentenceRef)
output := make(chan bucket.Bucket, 10)
go func() {
defer close(input)
for i, text := range sentences {
input <- bucket.SentenceRef{Sentence: text, Reference: i}
}
}()
go b.Run(input, output)
for bk := range output {
// Process bk.Batch, bk.Shape, bk.References...
}
func (*Bucketizer) ByPower ¶
func (b *Bucketizer) ByPower(batchSize, minSentenceLength int, base float64) *Bucketizer
ByPower configures the bucketizer to group sentences in buckets into fixed batchSize, where the sentence length is rounded up to the next power of base.
The minSentenceLength is the smallest sentenceLength bucket.
See also ByPowerBudget, if you want to round to a fixed tokens budget.
For full control on the bucketing function, see WithShapeFn.
func (*Bucketizer) ByPowerBudget ¶
func (b *Bucketizer) ByPowerBudget(tokensBudget, minSentenceLength int, base float64) *Bucketizer
ByPowerBudget configures the bucketizer to group sentences in buckets with a fixed total number of tokens (the tokensBudget), and where the sentence length is rounded up to the next power of base.
That means the batchSize is adjusted to keep the total number of tokens in the bucket close to tokensBudget, and hopefully the downstream tasks run on +/- constant time.
For sentence lengths > tokenBudget, it simply uses batchSize = 1.
For full control on the bucketing function, see WithShapeFn.
func (*Bucketizer) ByTwoBitBucket ¶
func (b *Bucketizer) ByTwoBitBucket(batchSize, minSentenceLength int) *Bucketizer
ByTwoBitBucket configures the bucketizer to use buckets of sentence-length sized to the next value that can be represented with 2 bits. So: 1, 2, 3, 4, 6, 8, 12, 16, ...
This is a "2-bit semi-log bucketing", and each size is separated from the other by a factor of 1.5 or 1.333 alternatingly, on average, by a factor of 1.414 (sqrt(2)), but results in numbers that are "friendlier" for binary addressing (and memory pages, etc.).
For full control on the bucketing function, see WithShapeFn.
func (*Bucketizer) ByTwoBitBucketBudget ¶
func (b *Bucketizer) ByTwoBitBucketBudget(tokensBudget, minSentenceLength int) *Bucketizer
ByTwoBitBucketBudget configures the bucketizer to use buckets of sentence-length sized to the next value that can be represented with 2 bits. So: 1, 2, 3, 4, 6, 8, 12, 16, ...
This is a "2-bit semi-log bucketing", and each size is separated from the other by a factor of 1.5 or 1.333 alternatingly, on average, by a factor of 1.414 (sqrt(2)), but results in numbers that are "friendlier" for binary addressing (and memory pages, etc.).
For full control on the bucketing function, see WithShapeFn.
func (*Bucketizer) Run ¶
func (b *Bucketizer) Run(input <-chan SentenceRef, output chan<- Bucket)
Run starts the streaming of the sentences into buckets.
It loops reading from the input, processing it, and writing to the output. The loop terminates when the input channel is closed.
The order of the output is very likely not the same as the input, but one can reconciliate the order by using the Reference.
By default this parallelizes the tokenization of sentences, but it can be configured with MaxParallelization.
func (*Bucketizer) WithBatchPadding ¶
func (b *Bucketizer) WithBatchPadding(useBatchPadding bool) *Bucketizer
WithBatchPadding sets whether partially emitted batches should be padded with empty sentences to match the maximum BatchSize.
This can be useful with WithMaxDelay -- it can also be configured there -- and if processing batches, to use on the final batches, when flushing the last buffered sentences.
func (*Bucketizer) WithMaxDelay ¶
func (b *Bucketizer) WithMaxDelay(d time.Duration, useBatchPadding bool) *Bucketizer
WithMaxDelay sets the maximum time a bucket can wait for more sentences before it is emitted partially.
A value of 0 (default) means no timeout—buckets only emit when full or closed.
An incomplete batch can either be returned partially (useBatchPadding == false) or padded with empty sentences (useBatchPadding == true). This is important if used fixed-shapes backends (e.g. XLA), which may expect full batches.
func (*Bucketizer) WithMaxParallelization ¶
func (b *Bucketizer) WithMaxParallelization(maxParallelization int) *Bucketizer
WithMaxParallelization sets the maximum number of sentences to tokenize in parallel.
Set to -1 to use runtime.NumCPU().
func (*Bucketizer) WithShapeFn ¶
func (b *Bucketizer) WithShapeFn(shapeFn ShapeFn) *Bucketizer
WithShapeFn sets the shape function to be used by the bucketizer.
See also ByPower and ByPowerBudget for common shape functions.
type Reference ¶
type Reference any
Reference is an opaque type that is to conciliate sentences streamed into input -- since their order is lost during the bucketing.
type SentenceRef ¶
SentenceRef pair are the inputs to be tokenized and bucketed.
Source Files