-
Notifications
You must be signed in to change notification settings - Fork 23
/
util.go
221 lines (179 loc) · 4.2 KB
/
util.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
package tokenizer
import (
"errors"
"log"
)
type TruncationParams struct {
MaxLength int
Strategy TruncationStrategy
Stride int
}
type PaddingParams struct {
Strategy PaddingStrategy
Direction PaddingDirection
PadId int
PadTypeId int
PadToken string
}
// PaddingStrategy is a enum of either
// - string `BatchLongest`
// - or a func type `Fixed(uint)` which return a uint
// Example:
//
// func main() {
// var ps PaddingStrategy
// ps = NewPaddingStrategy(WithFixed(3))
// fmt.Println(ps.Value)
// }
type PaddingStrategy struct {
Value interface{}
Name string
}
type PaddingStrategyOption func(*PaddingStrategy)
func WithBatchLongest() PaddingStrategyOption {
return func(ps *PaddingStrategy) {
ps.Value = "BatchLongest"
ps.Name = "BatchLongest"
}
}
func WithFixed(size int) PaddingStrategyOption {
return func(ps *PaddingStrategy) {
ps.Value = size
ps.Name = "Fixed"
}
}
func NewPaddingStrategy(opts ...PaddingStrategyOption) *PaddingStrategy {
const defaultVal = "BatchLongest"
ps := &PaddingStrategy{
Value: defaultVal,
Name: defaultVal,
}
for _, opt := range opts {
opt(ps)
}
return ps
}
// TruncationStrategy is enum of int type represents truncation strategy
type TruncationStrategy int
const (
LongestFirst TruncationStrategy = iota
OnlyFirst
OnlySecond
)
const (
SecondSequenceNotProvided = "Truncation error: Second sequence not provided"
SequenceTooShort = "Truncation error: Sequence to truncate too short to respect the provided max_length"
)
func TruncateEncodings(encoding, pairEncoding *Encoding, params *TruncationParams) (tEncoding, tPairEncoding *Encoding) {
var (
totalLength int
toRemove int
err error
)
if params.MaxLength == 0 {
return encoding, pairEncoding
}
totalLength = len(encoding.GetIds())
if pairEncoding != nil {
totalLength = len(encoding.GetIds()) + len(pairEncoding.GetIds())
}
if totalLength < params.MaxLength {
return encoding, pairEncoding
}
toRemove = totalLength - params.MaxLength
switch params.Strategy {
case LongestFirst:
nFirst := len(encoding.GetIds())
nSecond := 0
if pairEncoding != nil {
nSecond = len(pairEncoding.GetIds())
}
for i := 0; i < toRemove; i++ {
if nFirst > nSecond {
nFirst -= 1
}
nSecond -= 1
}
encoding.Truncate(nFirst, params.Stride)
if pairEncoding != nil {
pairEncoding.Truncate(nSecond, params.Stride)
}
case OnlyFirst, OnlySecond:
var truncateFunc = func(target *Encoding) (*Encoding, error) {
targetLength := len(target.GetIds())
if targetLength > toRemove {
target.Truncate(targetLength-toRemove, params.Stride)
return target, nil
} else {
err := errors.New(SequenceTooShort)
return nil, err
}
}
if params.Strategy == OnlyFirst {
encoding, err = truncateFunc(encoding)
} else if pairEncoding != nil {
pairEncoding, err = truncateFunc(pairEncoding)
} else {
err = errors.New(SecondSequenceNotProvided)
}
}
if err != nil {
log.Fatal(err)
}
return encoding, pairEncoding
}
func PadEncodings(encodings []Encoding, params PaddingParams) []Encoding {
if len(encodings) == 0 {
return encodings
}
var padLength int
switch params.Strategy.Name {
case "Fixed":
padLength = params.Strategy.Value.(int)
case "BatchLongest":
var max int = 0
for _, encoding := range encodings {
if len(encoding.GetIds()) > max {
max = len(encoding.GetIds())
}
}
padLength = max
}
// TODO: implement concurrency with for loop
var newEncodings []Encoding
for _, e := range encodings {
en := e
paddedEn := en.Pad(padLength, params.PadId, params.PadTypeId, params.PadToken, params.Direction)
newEncodings = append(newEncodings, *paddedEn)
}
return newEncodings
}
type Range []int
func NewRange(start, end int) Range {
if start < 0 {
panic("Invalid 'start' for NewRange()")
}
if end < 0 || end <= start {
panic("Invalid 'end' for NewRange()")
}
var r []int
for i := start; i < end; i++ {
r = append(r, i)
}
return r
}
func (r Range) Len() int {
return len(r)
}
func (r Range) Contains(item int) bool {
for _, v := range r {
if v == item {
return true
}
}
return false
}
func (r Range) IsEmpty() bool {
return len(r) == 0
}
// TODO. more methods of Range