// Portions of this file are derived from github.com/ikawaha/kagome
// Copyright (c) 2014-2023 Ikawaha
// Licensed under the MIT License.
package morph
import (
"encoding/binary"
"unicode/utf8"
)
type nodeClass uint8
const (
classDummy nodeClass = iota
classKnown
classUnknown
)
type latticeNode struct {
id int32
position int // byte offset in input
start int // rune offset in input
class nodeClass
cost int32
left int16
right int16
weight int16
surface string
prev int32 // index in nodes slice, -1 for none
}
type lattice struct {
nodes []latticeNode
list [][]int32
path []int32
}
func newLattice() *lattice {
return &lattice{
nodes: make([]latticeNode, 0, 256),
list: make([][]int32, 0, 64),
path: make([]int32, 0, 32),
}
}
// Tokenize segments input string into IPADIC morphemes matching kagome Normal mode.
func (d *Dict) Tokenize(inp string) []Token {
if len(inp) == 0 {
return nil
}
rc := utf8.RuneCountInString(inp)
la := d.pool.Get().(*lattice)
la.nodes = la.nodes[:0]
if cap(la.list) < rc+2 {
la.list = make([][]int32, rc+2)
} else {
la.list = la.list[:rc+2]
for i := range la.list {
la.list[i] = la.list[i][:0]
}
}
la.path = la.path[:0]
addNode := func(runePos, id, position, start int, class nodeClass, surface string) {
var left, right, weight int16
switch class {
case classKnown:
off := id * 6
left = int16(binary.LittleEndian.Uint16(d.morphBytes[off : off+2]))
right = int16(binary.LittleEndian.Uint16(d.morphBytes[off+2 : off+4]))
weight = int16(binary.LittleEndian.Uint16(d.morphBytes[off+4 : off+6]))
case classUnknown:
off := id * 6
left = int16(binary.LittleEndian.Uint16(d.unkMorphs[off : off+2]))
right = int16(binary.LittleEndian.Uint16(d.unkMorphs[off+2 : off+4]))
weight = int16(binary.LittleEndian.Uint16(d.unkMorphs[off+4 : off+6]))
}
nodeIdx := int32(len(la.nodes))
la.nodes = append(la.nodes, latticeNode{
id: int32(id),
position: position,
start: start,
class: class,
left: left,
right: right,
weight: weight,
surface: surface,
prev: -1,
})
p := runePos + utf8.RuneCountInString(surface)
la.list[p] = append(la.list[p], nodeIdx)
}
// BOS (rune 0)
addNode(0, -1, 0, 0, classDummy, inp[0:0])
// EOS (rune rc+1)
addNode(rc+1, -1, len(inp), rc, classDummy, inp[len(inp):len(inp)])
runePos := -1
for pos, ch := range inp {
runePos++
anyMatches := false
// Known words
d.commonPrefixSearchCallback(inp[pos:], func(id, l int) {
addNode(runePos, id, pos, runePos, classKnown, inp[pos:pos+l])
anyMatches = true
})
// Unknown words
class := d.charCategory(ch)
ccInfo := &d.charClasses[class]
if !anyMatches || ccInfo.Invoke {
var endPos int
if ch != utf8.RuneError {
endPos = pos + utf8.RuneLen(ch)
} else {
endPos = pos + 1
}
unkWordLen := 1
if ccInfo.Group {
for i, w, size := endPos, 0, len(inp); i < size; i += w {
var c rune
c, w = utf8.DecodeRuneInString(inp[i:])
if d.charCategory(c) != class {
break
}
endPos += w
unkWordLen++
if unkWordLen >= 1024 {
break
}
}
}
prev := pos
if c, size := utf8.DecodeLastRuneInString(inp[pos:endPos]); c != utf8.RuneError {
prev = endPos - size
}
id := ccInfo.UnkIndex
dup := ccInfo.UnkDup
for x := range int(dup) + 1 {
if pos < prev {
addNode(runePos, int(id)+x, pos, runePos, classUnknown, inp[pos:prev])
}
addNode(runePos, int(id)+x, pos, runePos, classUnknown, inp[pos:endPos])
}
}
}
// Forward pass (Viterbi)
const maxCost = 1<<31 - 1
for i := 1; i < len(la.list); i++ {
currentList := la.list[i]
for _, targetIdx := range currentList {
target := &la.nodes[targetIdx]
prevList := la.list[target.start]
if len(prevList) == 0 {
target.cost = maxCost
continue
}
for j, prevIdx := range prevList {
n := &la.nodes[prevIdx]
c := d.connCost(n.right, target.left)
totalCost := int64(c) + int64(target.weight) + int64(n.cost)
if totalCost > maxCost {
totalCost = maxCost
}
if j == 0 || int32(totalCost) < target.cost {
target.cost = int32(totalCost)
target.prev = prevIdx
}
}
}
}
// Backward pass
if len(la.list[rc+1]) == 0 {
if cap(la.nodes) <= 16384 {
d.pool.Put(la)
}
return nil
}
eosIdx := la.list[rc+1][0]
p := la.nodes[eosIdx].prev
for p != -1 {
node := &la.nodes[p]
if node.id == -1 { // BOS reached
break
}
la.path = append(la.path, p)
p = node.prev
}
tokens := make([]Token, len(la.path))
for i := range la.path {
node := &la.nodes[la.path[len(la.path)-1-i]]
tok := Token{
Surface: node.surface,
Start: node.position,
End: node.position + len(node.surface),
}
if node.class == classKnown {
tok.Known = true
featOff := uint32(node.id) * 16
posID := d.featBytes[featOff]
itID := d.featBytes[featOff+1]
ifID := d.featBytes[featOff+2]
baseOff := binary.LittleEndian.Uint32(d.featBytes[featOff+4 : featOff+8])
readOff := binary.LittleEndian.Uint32(d.featBytes[featOff+8 : featOff+12])
pronOff := binary.LittleEndian.Uint32(d.featBytes[featOff+12 : featOff+16])
tok.POS = d.posTuples[posID]
tok.InflType = d.inflTypes[itID]
tok.InflForm = d.inflForms[ifID]
tok.Base = d.stringAtOrSurface(baseOff, tok.Surface)
tok.Reading = d.stringAt(readOff)
if pronOff == readOff {
tok.Pron = tok.Reading
} else {
tok.Pron = d.stringAt(pronOff)
}
} else {
tok.Known = false
tok.POS = d.unkFeatures[node.id]
tok.Base = tok.Surface
}
tokens[i] = tok
}
if cap(la.nodes) <= 16384 {
d.pool.Put(la)
}
return tokens
}