tartrazine/classifier.go

103 lines
2.7 KiB
Go
Raw Normal View History

2017-06-13 11:56:07 +00:00
package enry
import (
"math"
2017-06-12 11:42:20 +00:00
"sort"
2017-06-08 07:27:27 +00:00
"gopkg.in/src-d/enry.v1/internal/tokenizer"
)
2017-06-12 11:42:20 +00:00
// Classifier is the interface in charge to detect the possible languages of the given content based on a set of
// candidates. Candidates is a map which can be used to assign weights to languages dynamically.
type Classifier interface {
2017-06-12 11:42:20 +00:00
Classify(content []byte, candidates map[string]float64) (languages []string)
}
type classifier struct {
languagesLogProbabilities map[string]float64
tokensLogProbabilities map[string]map[string]float64
tokensTotal float64
}
2017-06-12 11:42:20 +00:00
type scoredLanguage struct {
language string
score float64
}
// Classify returns a sorted slice of possible languages sorted by decreasing language's probability
func (c *classifier) Classify(content []byte, candidates map[string]float64) []string {
if len(content) == 0 {
return nil
}
2017-05-31 10:07:46 +00:00
var languages map[string]float64
if len(candidates) == 0 {
languages = c.knownLangs()
} else {
2017-05-31 10:07:46 +00:00
languages = make(map[string]float64, len(candidates))
for candidate, weight := range candidates {
if lang, ok := GetLanguageByAlias(candidate); ok {
2017-06-13 11:56:07 +00:00
candidate = lang
}
2017-06-13 11:56:07 +00:00
languages[candidate] = weight
}
}
tokens := tokenizer.Tokenize(content)
2017-06-12 11:42:20 +00:00
scoredLangs := make([]*scoredLanguage, 0, len(languages))
2017-05-31 10:07:46 +00:00
for language := range languages {
2017-06-12 11:42:20 +00:00
scoredLang := &scoredLanguage{
language: language,
score: c.tokensLogProbability(tokens, language) + c.languagesLogProbabilities[language],
}
scoredLangs = append(scoredLangs, scoredLang)
}
return sortLanguagesByScore(scoredLangs)
}
func sortLanguagesByScore(scoredLangs []*scoredLanguage) []string {
sort.Stable(byScore(scoredLangs))
2017-06-12 11:42:20 +00:00
sortedLanguages := make([]string, 0, len(scoredLangs))
for _, scoredLang := range scoredLangs {
sortedLanguages = append(sortedLanguages, scoredLang.language)
}
2017-06-12 11:42:20 +00:00
return sortedLanguages
}
2017-05-31 10:07:46 +00:00
func (c *classifier) knownLangs() map[string]float64 {
langs := make(map[string]float64, len(c.languagesLogProbabilities))
for lang := range c.languagesLogProbabilities {
2017-05-31 10:07:46 +00:00
langs[lang]++
}
return langs
}
func (c *classifier) tokensLogProbability(tokens []string, language string) float64 {
var sum float64
for _, token := range tokens {
sum += c.tokenProbability(token, language)
}
return sum
}
func (c *classifier) tokenProbability(token, language string) float64 {
tokenProb, ok := c.tokensLogProbabilities[language][token]
if !ok {
tokenProb = math.Log(1.000000 / c.tokensTotal)
}
return tokenProb
}
type byScore []*scoredLanguage
func (b byScore) Len() int { return len(b) }
func (b byScore) Swap(i, j int) { b[i], b[j] = b[j], b[i] }
func (b byScore) Less(i, j int) bool { return b[j].score < b[i].score }