【Chinese Word Segmentation Series】 1. Fast Segmentation Based on the Aho-Corasick Automaton

【中文分词系列】 1. 基于AC自动机的快速分词 苏剑林 · 2016-08-17 · Machine Learning dynamic programming AC automaton tokenization NLP vocabulary

Foreword: I spent quite a lot of time this summer on Chinese word segmentation and language models, hit countless walls, and picked up a few scattered insights along the way. I plan to write a series to share what I've learned. Although I call it a "series," it's really just a collection of notes rather than a systematic tutorial, so please bear with me.

Chinese Word Segmentation

I won't dwell too much on the background and importance of Chinese word segmentation — matrix67's post here gives a very clear introduction to segmentation and segmentation algorithms, and it's worth a read. In text mining, although quite a few articles have already explored approaches that skip segmentation entirely, such as this blog's own Text Sentiment Classification (III): To Segment or Not to Segment, in most cases segmentation is still treated as the first step of text mining. So an effective segmentation algorithm remains important. Of course, as the first step, Chinese word segmentation has already been studied for a very long time, and a lot of the current work is more summary in nature — at best a minor improvement, rather than any major breakthrough.

Currently there are two main approaches to Chinese word segmentation: dictionary lookup and character tagging. Dictionary lookup methods include: the mechanical maximum matching method, the minimum-word-count method, maximum probability combination based on a directed acyclic graph, and maximum probability combination based on a language model, among others. Dictionary lookup methods are simple and efficient (thanks to dynamic programming), and in particular the maximum probability method combined with a language model can handle ambiguity quite well. However, it cannot solve one of the two major difficulties in Chinese segmentation — out-of-vocabulary words (the two major difficulties in Chinese segmentation being ambiguity and out-of-vocabulary words). To address this, people have also proposed the character-tagging approach: so-called character tagging uses a handful of tags (for example, in the 4-tag scheme: single, a single character forming its own word; begin, the start of a multi-character word; middle, the middle part of a word with three or more characters; end, the end of a multi-character word) to represent the correct segmentation of a sentence. This turns segmentation into a sequence-to-sequence problem (input sentence to tag sequence), which handles out-of-vocabulary words fairly well, but is slower, and in scenarios where a complete dictionary is already available, character tagging may actually perform worse than dictionary lookup. In short, each has its pros and cons (which sounds like a truism), and in practice the two are often combined — for instance, Jieba uses maximum probability combination over a directed acyclic graph, while for consecutive single characters it falls back to an HMM model based on character tagging to identify them. more

The Aho-Corasick Automaton

What we'll implement first in this post is dictionary-lookup-based segmentation. The dictionary lookup process is: 1. Given a batch of words, check whether a given sentence contains any of them; 2. If it does, how do we resolve the resulting ambiguity? Step 1 is known in computer science as "multi-pattern matching." It looks simple, but implementing it efficiently is not trivial at all. A complete dictionary typically has well over a hundred thousand words, and if you search by enumerating each one individually, the computational cost becomes unbearable. In fact, that's not how humans do it either — when we look something up in a dictionary, we first check the initial letter, then narrow our search to entries sharing that initial letter, then compare the next letter, and so on. This requires two things: 1. A dictionary with a special ordering; 2. Effective search techniques. For the first requirement we have the so-called prefix tree (trie), and for the second we have some classic algorithms, such as the Aho-Corasick automaton.

I won't comment much further on these two ingredients — not because I don't want to, but because that's about the extent of my own understanding. As far as I understand it, the Aho-Corasick automaton is simply an efficient multi-pattern matching algorithm built on a trie data structure. I don't even need to implement it myself, since Python already has a library for it: pyahocorasick. So all we need to worry about is how to use it. The official tutorial already introduces the basic usage of pyahocorasick in detail, so I won't repeat that here. (Unfortunately, although pyahocorasick now supports both python2 and python3, in python2 it only supports byte strings and not unicode strings, whereas in python3 it defaults to unicode encoding, which can be a bit confusing when writing code — though it's not a fundamental issue. This post uses Python 2.7.)

To build an Aho-Corasick-based segmentation system, we first need a text dictionary — suppose the dictionary has two columns per line, the word and its corresponding frequency, separated by a space. Then we can build an Aho-Corasick automaton with the following code.

import ahocorasick

def load_dic(dicfile):
    from math import log
    dic = ahocorasick.Automaton()
    total = 0.0
    with open(dicfile) as dicfile:
        words = []
        for line in dicfile:
            line = line.split(' ')
            words.append((line[0], int(line[1])))
            total += int(line[1])    
    for i,j in words:
        dic.add_word(i, (i, log(j/total))) #这里使用了对数概率，防止溢出
    dic.make_automaton()
    return dic

dic = load_dic('me.dic')

One very user-friendly feature of pyahocorasick's automaton construction is that it lets you add vocabulary as "key–annotation" pairs (note the line dic.add_word(i, (i, log(j/total)))). This means we can attach whatever information we want in the annotation — frequency, part of speech, and so on — and it will be returned together with the match during lookup. With the automaton built above, we can easily implement a "full-mode" segmentation, i.e. scanning out every word in the dictionary that appears in the text (which is, after all, exactly what the automaton is built for).

def all_cut(sentence):
    words = []
    for i,j in dic.iter(sentence):
        words.append(j[0])
    return words

For a long sentence, this can return a huge number of matched words, so use it with caution.

Maximum Matching

Of course, the so-called full-mode segmentation above barely counts as segmentation at all — it's really just a lookup. Let's now implement a classic segmentation algorithm: maximum matching.

Maximum matching scans from left to right, matching against words in the dictionary, and always picks the longest matching word. This is a fairly crude segmentation method, as also mentioned in matrix67's article — it's easy to construct counterexamples: if the dictionary contains the words "不" (not), "不可" (cannot), "能" (able), and "可能" (possible), but not "不可能" (impossible), then "不可能" would incorrectly be split into "不可/能". Even so, when high precision isn't required, this segmentation method is still acceptable, since it's very fast. Below is an implementation of maximum matching based on the Aho-Corasick automaton:

def max_match_cut(sentence):
    sentence = sentence.decode('utf-8')
    words = ['']
    for i in sentence:
        i = i.encode('utf-8')
        if dic.match(words[-1] + i):
            words[-1] += i
        else:
            words.append(i)
    return words

The code is short and quite clear, mainly relying on pyahocorasick's match function. Testing on my machine, this algorithm runs at roughly 4M characters/s. According to the author of HanLP, doing something similar in Java can reach speeds of 20M characters/s! Doing this in Python faces two constraints: one is the inherent speed limitation of Python itself, and the other is a limitation of pyahocorasick — since it doesn't support unicode encoding, Chinese characters have variable byte lengths, requiring constant re-encoding just to determine the character length, which means the implementation above is not actually maximally efficient.

The maximum matching method described above is, strictly speaking, "forward maximum matching." Similarly, there's also "backward maximum matching," which, as the name suggests, scans the sentence from right to left performing maximum matching; this generally performs a bit better than forward maximum matching. If we want to implement it with the Aho-Corasick automaton, the only way is to store all dictionary words in reverse order, reverse the input sentence as well, and then perform forward maximum matching on the reversed sentence.

Maximum Probability Combination

The maximum probability combination method is currently one of the better approaches, balancing speed and accuracy. The idea is: for a given sentence, if splitting it into the words $w_1,w_2,\dots,w_n$ is the optimal segmentation, then it should maximize the following probability:

$$P(w_1,w_2,\dots,w_n)$$

Directly estimating this probability isn't easy, so some approximation scheme is generally used, for example

$$P(w_1,w_2,\dots,w_n)\approx P(w_1)P(w_2|w_1)P(w_3|w_2)\dots P(w_n|w_{n-1})$$

Here $P(w_k|w_{k-1})$ is what's called a language model, and it already accounts for semantics to some extent. Of course, it's very difficult for an ordinary segmentation tool to estimate $P(w_k|w_{k-1})$, so a simpler approximation is generally adopted instead.

$$P(w_1,w_2,\dots,w_n)\approx P(w_1)P(w_2)P(w_3)\dots P(w_n)$$

Viewed through the lens of graph theory, this is exactly the maximum probability path in a directed acyclic graph.

Below I show how to implement the latter scheme using the Aho-Corasick automaton combined with dynamic programming.

def max_proba_cut(sentence):
    paths = {0:([], 0)}
    end = 0
    for i,j in dic.iter(sentence):
        start,end = 1+i-len(j[0]), i+1
        if start not in paths:
            last = max([i for i in paths if i < start])
            paths[start] = (paths[last][0]+[sentence[last:start]], paths[last][1]-10)
        proba = paths[start][1]+j[1]
        if end not in paths or proba > paths[end][1]:
            paths[end] = (paths[start][0]+[j[0]], proba)
    if end < len(sentence):
        return paths[end][0] + [sentence[end:]]
    else:
        return paths[end][0]

The code is again short and clear. Here, the frequency for unmatched portions is assumed to be $e^{-10}$, which can be adjusted as needed. One thing to note is that, because of the different approach used, the dynamic programming scheme here differs from the usual directed-acyclic-graph-based dynamic programming, though the underlying idea is quite natural. Be aware that if you apply this function directly to a sentence with tens of thousands of characters, it will be fairly slow and memory-intensive, because a dictionary is used to keep track of every intermediate candidate during the dynamic programming process. Fortunately, Chinese sentences already have plenty of natural sentence-breaking markers, such as punctuation and line breaks. We can use these markers to split the sentence into smaller parts and segment them one by one, as shown below.

to_break = ahocorasick.Automaton()
for i in ['，', '。', '！', '、', '？', ' ', '\n']:
    to_break.add_word(i, i)

to_break.make_automaton()

def map_cut(sentence):
    start = 0
    words = []
    for i in to_break.iter(sentence):
        words.extend(max_proba_cut(sentence[start:i[0]+1]))
        start = i[0]+1
    words.extend(max_proba_cut(sentence[start:]))
    return words

On a server, I sampled 100,000 articles (over 100 million characters) and compared the speed against Jieba, using the same dictionary and with Jieba's new-word discovery disabled. I found that the map_cut segmentation implemented with the Aho-Corasick automaton was roughly 2–3 times faster than Jieba, running at about 1M characters/s.

Finally, it's worth mentioning that the implementation approach used in max_proba_cut can also be applied to other segmentation methods involving dynamic programming, such as minimum-word-count segmentation:

def min_words_cut(sentence):
    paths = {0:([], 0)}
    end = 0
    for i,j in dic.iter(sentence):
        start,end = 1+i-len(j[0]), i+1
        if start not in paths:
            last = max([i for i in paths if i < start])
            paths[start] = (paths[last][0]+[sentence[last:start]], paths[last][1]+1)
        num = paths[start][1]+1
        if end not in paths or num < paths[end][1]:
            paths[end] = (paths[start][0]+[j[0]], num)
    if end < len(sentence):
        return paths[end][0] + [sentence[end:]]
    else:
        return paths[end][0]

Here a penalty rule is adopted: each word incurs a penalty point, with an extra penalty point for out-of-vocabulary words, and the segmentation with the fewest total penalty points wins.

Summary

In fact, wherever dictionary lookup is involved, the Aho-Corasick automaton has some role to play. Applying it to word segmentation is, in fact, a very natural application. We look forward to seeing more data structures and algorithms with better support for Chinese emerge, which would allow us to design even more efficient algorithms.