Fix style

Signed-off-by: Irina Khismatullina <irenekhismatullina@gmail.com>

Author: irinakhismatullina

bpe.go

@@ -8,44 +8,50 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-type TokenId int32
+// TokenID is a numerical identitier of the subword token
+type TokenID uint32
 
-type DecodedToken []TokenId
+// EncodedToken is a sequence of subword tokens ids
+type EncodedToken []TokenID
 
-type Rule struct {
-	left   TokenId
-	right  TokenId
-	result TokenId
+type rule struct {
+	left   TokenID
+	right  TokenID
+	result TokenID
 }
 
-type SpecialTokens struct {
-	unk TokenId
-	pad TokenId
-	bos TokenId
-	eos TokenId
+type specialTokens struct {
+	unk int
+	pad int
+	bos int
+	eos int
 }
 
+// Model is a Byte-Pair encoding model, which supports encoding and decoding text into sequences
+// of most frequent subword tokens
 type Model struct {
-	char2id       map[rune]TokenId
-	id2char       map[TokenId]rune
-	rules         []Rule
-	recipe        map[TokenId]DecodedToken
-	revRecipe     map[string]TokenId
-	specialTokens SpecialTokens
+	char2id       map[rune]TokenID
+	id2char       map[TokenID]rune
+	rules         []rule
+	recipe        map[TokenID]EncodedToken
+	revRecipe     map[string]TokenID
+	specialTokens specialTokens
 }
 
-func NewModel(nRules int) *Model {
+func newModel(nRules int) *Model {
 	return &Model{
-		make(map[rune]TokenId),
-		make(map[TokenId]rune),
-		make([]Rule, nRules),
-		make(map[TokenId]DecodedToken),
-		make(map[string]TokenId),
-		SpecialTokens{-1, -1, -1, -1},
+		make(map[rune]TokenID),
+		make(map[TokenID]rune),
+		make([]rule, nRules),
+		make(map[TokenID]EncodedToken),
+		make(map[string]TokenID),
+		specialTokens{-1, -1, -1, -1},
 	}
 }
 
-func DecodedTokenToString(token DecodedToken, id2char map[TokenId]rune) (string, error) {
+// DecodeToken converts the sequence of chars' ids into the string -
+// sequence of the corresponding chars
+func DecodeToken(token EncodedToken, id2char map[TokenID]rune) (string, error) {
 	word := ""
 	for _, id := range token {
 		if char, ok := id2char[id]; ok {
@@ -57,7 +63,8 @@ func DecodedTokenToString(token DecodedToken, id2char map[TokenId]rune) (string,
 	return word, nil
 }
 
-func ReadModel(reader io.Reader) (*Model, error) {
+// ReadModelFromText loads the BPE model from the text dump
+func ReadModelFromText(reader io.Reader) (*Model, error) {
 	scanner := bufio.NewScanner(reader)
 	var nChars, nRules int
 	scanner.Scan()
@@ -66,24 +73,23 @@ func ReadModel(reader io.Reader) (*Model, error) {
 		logrus.Fatal("Wrong input format: ", err)
 		return &Model{}, err
 	}
-	model := NewModel(nRules)
-	model.rules = make([]Rule, nRules)
+	model := newModel(nRules)
 	for i := 0; i < nChars; i++ {
 		var char rune
-		var charId TokenId
+		var charID TokenID
 		scanner.Scan()
-		_, err = fmt.Sscanf(scanner.Text(), "%d %d", &char, &charId)
+		_, err = fmt.Sscanf(scanner.Text(), "%d %d", &char, &charID)
 		if err != nil {
 			logrus.Fatal("Wrong input format: ", err)
 			return model, err
 		}
-		model.char2id[char] = charId
-		model.id2char[charId] = char
-		model.recipe[charId] = DecodedToken{charId}
-		model.revRecipe[string(char)] = charId
+		model.char2id[char] = charID
+		model.id2char[charID] = char
+		model.recipe[charID] = EncodedToken{charID}
+		model.revRecipe[string(char)] = charID
 	}
 	for i := 0; i < nRules; i++ {
-		var rule Rule
+		var rule rule
 		scanner.Scan()
 		_, err = fmt.Sscanf(scanner.Text(), "%d %d %d", &rule.left, &rule.right, &rule.result)
 		if err != nil {
@@ -92,7 +98,7 @@ func ReadModel(reader io.Reader) (*Model, error) {
 		}
 		model.rules[i] = rule
 		model.recipe[rule.result] = append(model.recipe[rule.left], model.recipe[rule.right]...)
-		resultString, err := DecodedTokenToString(model.recipe[rule.result], model.id2char)
+		resultString, err := DecodeToken(model.recipe[rule.result], model.id2char)
 		if err != nil {
 			logrus.Fatal("Unexpected token id inside the rules: ", err)
 			return model, err

bpe_test.go

@@ -1,19 +1,20 @@
 package bpe
 
 import (
-	"github.com/stretchr/testify/require"
 	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestNewModel(t *testing.T) {
-	model := NewModel(10)
+	model := newModel(10)
 	require.Equal(t, 10, len(model.rules))
 }
 
 func TestDecodedTokenToString(t *testing.T) {
-	id2char := map[TokenId]rune{1: []rune("a")[0], 2: []rune("b")[0], 3: []rune("c")[0]}
-	word, err := DecodedTokenToString(DecodedToken{1, 2, 1, 3, 3}, id2char)
+	id2char := map[TokenID]rune{1: []rune("a")[0], 2: []rune("b")[0], 3: []rune("c")[0]}
+	word, err := DecodeToken(EncodedToken{1, 2, 1, 3, 3}, id2char)
 	require.NoError(t, err)
 	require.Equal(t, "abacc", word)
 }
@@ -31,14 +32,14 @@ func TestReadModel(t *testing.T) {
 4 7 12
 1 0 2 4`)
 	expected := Model{
-		map[rune]TokenId{97: 8, 98: 7, 99: 6, 100: 5, 95: 4},
-		map[TokenId]rune{4: 95, 5: 100, 6: 99, 7: 98, 8: 97},
-		[]Rule{{4, 8, 9}, {4, 6, 10}, {4, 5, 11}, {4, 7, 12}},
-		map[TokenId]DecodedToken{4: {4}, 5: {5}, 6: {6}, 7: {7}, 8: {8}, 9: {4, 8}, 10: {4, 6}, 11: {4, 5}, 12: {4, 7}},
-		map[string]TokenId{"a": 8, "b": 7, "c": 6, "d": 5, "_": 4,
+		map[rune]TokenID{97: 8, 98: 7, 99: 6, 100: 5, 95: 4},
+		map[TokenID]rune{4: 95, 5: 100, 6: 99, 7: 98, 8: 97},
+		[]rule{{4, 8, 9}, {4, 6, 10}, {4, 5, 11}, {4, 7, 12}},
+		map[TokenID]EncodedToken{4: {4}, 5: {5}, 6: {6}, 7: {7}, 8: {8}, 9: {4, 8}, 10: {4, 6}, 11: {4, 5}, 12: {4, 7}},
+		map[string]TokenID{"a": 8, "b": 7, "c": 6, "d": 5, "_": 4,
 			"_a": 9, "_b": 12, "_c": 10, "_d": 11},
-		SpecialTokens{1, 0, 2, 4},
+		specialTokens{1, 0, 2, 4},
 	}
-	model, _ := ReadModel(reader)
+	model, _ := ReadModelFromText(reader)
 	require.Equal(t, expected, *model)
 }