Moved over from lingo, and tightened up the API surface

Author: chewxy

consopt.go

@@ -0,0 +1,152 @@
+package corpus
+
+import (
+	"log"
+	"sort"
+	"sync/atomic"
+	"unicode/utf8"
+
+	"github.com/pkg/errors"
+	"github.com/xtgo/set"
+)
+
+// ConsOpt is a construction option for manual creation of a Corpus
+type ConsOpt func(c *Corpus) error
+
+// WithWords creates a corpus from a word list. It may have repeated words
+func WithWords(a []string) ConsOpt {
+	f := func(c *Corpus) error {
+		s := set.Strings(a)
+		c.words = s
+		c.frequencies = make([]int, len(s))
+
+		ids := make(map[string]int)
+		maxID := len(s)
+
+		var totalFreq, maxWL int
+		// NOTE: here we're iterating over the set of words
+		for i, w := range s {
+			runeCount := utf8.RuneCountInString(w)
+			if runeCount > c.maxWordLength {
+				maxWL = runeCount
+			}
+
+			ids[w] = i
+		}
+
+		// NOTE: here we're iterating over the original word list.
+		for _, w := range a {
+			c.frequencies[ids[w]]++
+			totalFreq++
+		}
+
+		c.ids = ids
+		atomic.AddInt64(&c.maxid, int64(maxID))
+		c.totalFreq = totalFreq
+		c.maxWordLength = maxWL
+		return nil
+	}
+	return f
+}
+
+// WithOrderedWords creates a Corpus with the given word order
+func WithOrderedWords(a []string) ConsOpt {
+	f := func(c *Corpus) error {
+		s := a
+		c.words = s
+		c.frequencies = make([]int, len(s))
+		for i := range c.frequencies {
+			c.frequencies[i] = 1
+		}
+
+		ids := make(map[string]int)
+		maxID := len(s)
+		totalFreq := len(s)
+		var maxWL int
+		for i, w := range a {
+			runeCount := utf8.RuneCountInString(w)
+			if runeCount > c.maxWordLength {
+				maxWL = runeCount
+			}
+			ids[w] = i
+		}
+
+		c.ids = ids
+		atomic.AddInt64(&c.maxid, int64(maxID))
+		c.totalFreq = totalFreq
+		c.maxWordLength = maxWL
+		return nil
+	}
+	return f
+}
+
+// WithSize preallocates all the things in Corpus
+func WithSize(size int) ConsOpt {
+	return func(c *Corpus) error {
+		c.words = make([]string, 0, size)
+		c.frequencies = make([]int, 0, size)
+		return nil
+	}
+}
+
+// FromDict is a construction option to take a map[string]int where the int represents the word ID.
+// This is useful for constructing corpuses from foreign sources where the ID mappings are important
+func FromDict(d map[string]int) ConsOpt {
+	return func(c *Corpus) error {
+		var a sortutil
+		for k, v := range d {
+			a.words = append(a.words, k)
+			a.ids = append(a.ids, v)
+		}
+		sort.Sort(&a)
+		c.ids = make(map[string]int)
+		for i, w := range a.words {
+			if i != a.ids[i] {
+				return errors.Errorf("Unmarshaling error. Expected %dth ID to be %d. Got %d instead. Perhaps something went wrong during sorting? SLYTHERIN IT IS!", i, i, a.ids[i])
+			}
+			c.words = append(c.words, w)
+			c.frequencies = append(c.frequencies, 1)
+			c.ids[w] = i
+
+			c.totalFreq++
+			runeCount := utf8.RuneCountInString(w)
+			if runeCount > c.maxWordLength {
+				log.Printf("FD MaxWordLength %d - %q", runeCount, w)
+				c.maxWordLength = runeCount
+			}
+		}
+		c.maxid = int64(len(a.words))
+		return nil
+	}
+
+}
+
+// FromDictWithFreq is like FromDict, but also has a frequency.
+func FromDictWithFreq(d map[string]struct{ ID, Freq int }) ConsOpt {
+	return func(c *Corpus) error {
+		var a sortutil
+		for k, v := range d {
+			a.words = append(a.words, k)
+			a.ids = append(a.ids, v.ID)
+			a.freqs = append(a.freqs, v.Freq)
+		}
+		sort.Sort(&a)
+		c.ids = make(map[string]int)
+		for i, w := range a.words {
+			if i != a.ids[i] {
+				return errors.Errorf("Unmarshaling error. Expected %dth ID to be %d. Got %d instead. Perhaps something went wrong during sorting? SLYTHERIN IT IS!", i, i, a.ids[i])
+			}
+			c.words = append(c.words, w)
+			c.frequencies = append(c.frequencies, a.freqs[i])
+			c.ids[w] = i
+
+			c.totalFreq += a.freqs[i]
+			runeCount := utf8.RuneCountInString(w)
+			if runeCount > c.maxWordLength {
+				c.maxWordLength = runeCount
+			}
+		}
+		c.maxid = int64(len(a.words))
+		return nil
+	}
+}

consts.go

@@ -0,0 +1,59 @@
+package corpus
+
+// NumberWords was generated with this python code
+/*
+	numberWords = {}
+
+	simple = '''zero one two three four five six seven eight nine ten eleven twelve
+	        thirteen fourteen fifteen sixteen seventeen eighteen nineteen
+	        twenty'''.split()
+	for i, word in zip(xrange(0, 20+1), simple):
+	    numberWords[word] = i
+
+	tense = '''thirty forty fifty sixty seventy eighty ninety hundred'''.split()
+	for i, word in zip(xrange(30, 100+1, 10), tense):
+		numberWords[word] = i
+
+	larges = '''thousand million billion trillion quadrillion quintillion sextillion septillion'''.split()
+	for i, word in zip(xrange(3, 24+1, 3), larges):
+		numberWords[word] = 10**i
+*/
+var NumberWords = map[string]int{
+	"zero":        0,
+	"one":         1,
+	"two":         2,
+	"three":       3,
+	"four":        4,
+	"five":        5,
+	"six":         6,
+	"seven":       7,
+	"eight":       8,
+	"nine":        9,
+	"ten":         10,
+	"eleven":      11,
+	"twelve":      12,
+	"thirteen":    13,
+	"fourteen":    14,
+	"fifteen":     15,
+	"sixteen":     16,
+	"nineteen":    19,
+	"seventeen":   17,
+	"eighteen":    18,
+	"twenty":      20,
+	"thirty":      30,
+	"forty":       40,
+	"fifty":       50,
+	"sixty":       60,
+	"seventy":     70,
+	"eighty":      80,
+	"ninety":      90,
+	"hundred":     100,
+	"thousand":    1000,
+	"million":     1000000,
+	"billion":     1000000000,
+	"trillion":    1000000000000,
+	"quadrillion": 1000000000000000,
+	// "quintillion": 1000000000000000000,
+	// "sextillion": 1000000000000000000000,
+	// "septillion": 1000000000000000000000000,
+}