feat: implement Analyze

kiwi.go

@@ -6,24 +6,24 @@ package kiwi
 */
 import "C"
 
-// KiwiBuildOption is a bitwise OR of the KiwiBuildOption values.
-type KiwiBuildOption int
+// BuildOption is a bitwise OR of the KiwiBuildOption values.
+type BuildOption int
 
 const (
-	KIWI_BUILD_LOAD_DEFAULT_DICT   KiwiBuildOption = 1
-	KIWI_BUILD_INTEGRATE_ALLOMORPH KiwiBuildOption = 2
-	KIWI_BUILD_DEFAULT             KiwiBuildOption = 3
+	KIWI_BUILD_LOAD_DEFAULT_DICT   BuildOption = C.KIWI_BUILD_LOAD_DEFAULT_DICT
+	KIWI_BUILD_INTEGRATE_ALLOMORPH BuildOption = C.KIWI_BUILD_INTEGRATE_ALLOMORPH
+	KIWI_BUILD_DEFAULT             BuildOption = C.KIWI_BUILD_DEFAULT
 )
 
-// KiwiAnalyzeOption is a bitwise OR of the KiwiAnalyzeOption values.
-type KiwiAnalyzeOption int
+// AnalyzeOption is a bitwise OR of the KiwiAnalyzeOption values.
+type AnalyzeOption int
 
 const (
-	KIWI_MATCH_URL     KiwiAnalyzeOption = 1
-	KIWI_MATCH_EMAIL   KiwiAnalyzeOption = 2
-	KIWI_MATCH_HASHTAG KiwiAnalyzeOption = 4
-	KIWI_MATCH_MENTION KiwiAnalyzeOption = 8
-	KIWI_MATCH_ALL     KiwiAnalyzeOption = 15
+	KIWI_MATCH_URL     AnalyzeOption = C.KIWI_MATCH_URL
+	KIWI_MATCH_EMAIL   AnalyzeOption = C.KIWI_MATCH_EMAIL
+	KIWI_MATCH_HASHTAG AnalyzeOption = C.KIWI_MATCH_HASHTAG
+	KIWI_MATCH_MENTION AnalyzeOption = C.KIWI_MATCH_MENTION
+	KIWI_MATCH_ALL     AnalyzeOption = C.KIWI_MATCH_ALL
 )
 
 // KiwiVersion returns the version of the kiwi library.
@@ -37,33 +37,65 @@ type Kiwi struct {
 }
 
 // New returns a new Kiwi instance.
-func New(modelPath string, numThread int, options KiwiBuildOption) *Kiwi {
+// Don't forget to call Close after this.
+func New(modelPath string, numThread int, options BuildOption) *Kiwi {
 	return &Kiwi{
 		handler: C.kiwi_init(C.CString(modelPath), C.int(numThread), C.int(options)),
 	}
 }
 
 // TokenInfo returns the token info for the given token(Str).
 type TokenInfo struct {
-	Str         string
-	Position    int
-	Length      int
-	WorPosition int
+	// Position is the index of this token appears in the original text.
+	Position int
+
+	// Tag represents a type of this token (e.g. VV, NNG, ...).
+	// TODO: convert string to enum
+	Tag      string
+
+	// Form is the actual string of this token.
+	Form     string
 }
 
-// KiwiTokenResult is a result for Analyze.
-type KiwiTokenResult struct {
+// TokenResult is a result for Analyze.
+type TokenResult struct {
 	Tokens []TokenInfo
-	Score  float64
+	Score  float32
 }
 
 // Analyze returns the result of the analysis.
-func (k *Kiwi) Analyze(text string, topN int, options KiwiAnalyzeOption) KiwiTokenResult {
+func (k *Kiwi) Analyze(text string, topN int, options AnalyzeOption) []TokenResult {
 	kiwiResH := C.kiwi_analyze(k.handler, C.CString(text), C.int(topN), C.int(options))
 
-	score := float64(C.kiwi_res_prob(kiwiResH, 0))
+	defer C.kiwi_res_close(kiwiResH)
+
+	resSize := int(C.kiwi_res_size(kiwiResH))
+	res := make([]TokenResult, resSize)
+
+	for i := 0; i < resSize; i++ {
+		tokens := make([]TokenInfo, int(C.kiwi_res_word_num(kiwiResH, C.int(i))))
 
-	return KiwiTokenResult{
-		Score: score,
+		for j := 0; j < len(tokens); j++ {
+			tokens[j] = TokenInfo{
+				Form:     C.GoString(C.kiwi_res_form(kiwiResH, C.int(i), C.int(j))),
+				Tag:      C.GoString(C.kiwi_res_tag(kiwiResH, C.int(i), C.int(j))),
+				Position: int(C.kiwi_res_position(kiwiResH, C.int(i), C.int(j))),
+			}
+		}
+
+		res[i] = TokenResult{
+			Tokens: tokens,
+			Score:  float32(C.kiwi_res_prob(kiwiResH, C.int(i))),
+		}
 	}
+
+	return res
+}
+
+// Close frees the resource allocated for Kiwi and returns the exit status.
+// This must be called after New.
+//
+// Returns 0 if successful.
+func (k *Kiwi) Close() int {
+	return int(C.kiwi_close(k.handler))
 }

kiwi_test.go

@@ -13,5 +13,50 @@ func TestKiwiVersion(t *testing.T) {
 func TestAnalyze(t *testing.T) {
 	kiwi := New("./ModelGenerator", 1, KIWI_BUILD_DEFAULT)
 	res := kiwi.Analyze("아버지가 방에 들어가신다", 1, KIWI_MATCH_ALL)
-	assert.NotEqual(t, 0, res.Score)
+
+	expected := []TokenResult{
+		{
+			Tokens: []TokenInfo{
+				{
+					Position: 0,
+					Tag:      "NNG",
+					Form:     "아버지",
+				},
+				{
+					Position: 3,
+					Tag:      "JKS",
+					Form:     "가",
+				},
+				{
+					Position: 5,
+					Tag:      "NNG",
+					Form:     "방",
+				},
+				{
+					Position: 6,
+					Tag:      "JKB",
+					Form:     "에",
+				},
+				{
+					Position: 8,
+					Tag:      "VV",
+					Form:     "들어가",
+				},
+				{
+					Position: 11,
+					Tag:      "EP",
+					Form:     "시",
+				},
+				{
+					Position: 12,
+					Tag:      "EF",
+					Form:     "ᆫ다",
+				},
+			},
+			Score: -38.967132568359375,
+		},
+	}
+
+	assert.Equal(t, expected, res)
+	assert.Equal(t, 0, kiwi.Close())
 }