maker

#!/usr/bin/env python
# -*- coding:utf-8 -*-

import re
import sys
from optparse import OptionParser

try:
    unicode # python2
    def u(str): return str.decode("utf-8")
    def b(str): return str
    def fileformat(str): return str.encode("utf-8")
    pass
except: # python3
    def u(str): return str
    def b(str): return str.encode("utf-8")
    def fileformat(str): return str
    pass

categories = frozenset([u'UP1', u'UP2', u'UP3',
              u'BP1', u'BP2',
              u'UW1', u'UW2', u'UW3', u'UW4', u'UW5', u'UW6',
              u'BW1', u'BW2', u'BW3',
              u'TW1', u'TW2', u'TW3', u'TW4',
              u'UC1', u'UC2', u'UC3', u'UC4', u'UC5', u'UC6',
              u'BC1', u'BC2', u'BC3',
              u'TC1', u'TC2', u'TC3', u'TC4',
              u'UQ1', u'UQ2', u'UQ3',
              u'BQ1', u'BQ2', u'BQ3', u'BQ4',
              u'TQ1', u'TQ2', u'TQ3', u'TQ4',
])
re_key = re.compile(u'([a-zA-Z0-9]*):(.*)')
def load(f, scale = 10000):
    model = dict([category, {}] for category in categories)
    model[u'BIAS'] = 0
    for line in f:
        a = u(line).strip().split('\t')
        if len(a) == 2:
            m = re_key.match(a[0])
            category = m.group(1)
            key = m.group(2)
            value = int(float(a[1]) * scale + 0.5)
            if category not in categories:
                continue
            model[category][key] = value
        else:
            model['BIAS'] = int(float(a[0]) * scale + 0.5)
    return model

languages = {}
class language(object):
    def __init__(self, name, ext):
        self.name = name
        self.ext = ext
        self.execute = None

    def __call__(self, func):
        languages[self.name] = self
        self.execute = func
        return func

class javascript(object):
    @classmethod
    def tostr(cls, obj):
        items = []
        for key, value in obj.items():
            if isinstance(value, dict):
                value = cls.tostr(value)
            if not re.match('[a-zA-Z0-9]+', key):
                key = key.replace("'", "\\'")
                key = "'" + key + "'"
            items.append(u"{0}:{1}".format(key, value))
        items.sort()
        return '{' + ','.join(items) + '}'

    @language('javascript', 'js')
    def original(model):
        template = u(open('templates/javascript.js').read())
        template = template.replace(u'__MODEL__', javascript.tostr(model))
        return template

    @language('javascriptdp', 'dp.js')
    def dp(model):
        template = u(open('templates/javascript.dp.js').read())
        template = template.replace(u'__MODEL__', javascript.tostr(model))
        return template

@language('python', 'py')
def python(model):
    def tostr(obj):
        items = []
        for key, value in obj.items():
            if isinstance(value, dict):
                value = tostr(value)
            key = key.replace(u"'", u"\\'")
            key = u"u'" + key + u"'"
            items.append(u"{0}:{1}".format(key, value))
        items.sort()
        return u'{' + u','.join(items) + u'}'

    template = u(open('templates/python.py').read())
    template = template.replace(u'__MODEL__', tostr(model))
    return template

@language('perl', 'pm')
def perl(model):
    def tostr(obj):
        items = []
        for key, value in obj.items():
            key = key.replace("'", "\\'")
            key = "'" + key + "'"
            items.append(key + '=>' + str(value))
        items.sort()
        return '(' + ','.join(items) + ')'

    model_str = ''
    for key, value in model.items():
        if key == 'BIAS':
            model_str += 'my $BIAS = ' + str(value) + ';\n'
        else:
            model_str += 'my %' + key + ' = ' + tostr(value) + ';\n'
    template = u(open('templates/perl.pl').read())
    template = template.replace(u'__MODEL__', model_str)
    return template

@language('ruby', 'rb')
def ruby(model):
    def tostr(obj):
        items = []
        for key, value in obj.items():
            key = key.replace("'", "\\'")
            key = "'" + key + "'"
            items.append(key + '=>' + str(value))
        items.sort()
        return '{' + ','.join(items) + '}'

    model_str = ''
    for key, value in model.items():
        if key == 'BIAS':
            model_str += 'BIAS = ' + str(value) + ';\n'
        else:
            model_str += key + ' = ' + tostr(value) + ';\n'
    template = u(open('templates/ruby.rb').read())
    template = template.replace(u'__MODEL__', model_str)
    return template

@language('cpp', 'hpp')
def cpp(model):
    def ch2label(ch):
        if isinstance(ch, int):
            ch = chr(ch)
        if 'a' <= ch <='z' or 'A' <= ch <= 'Z' or '0' <= ch <= '9':
            return "'" + ch + "'"
        else:
            return "'\\x" + hex(ord(ch))[2:] + "'"
    def tostr(obj, pos = 0):
        if len(obj) == 1:
            key, value = list(obj.items())[0]
            conditions = ['s[' + str(i) + ']==' + ch2label(ch) for i, ch in enumerate(key.encode('utf-8'))][pos:]
            if len(conditions)==0:
                return 'return ' + str(value) + ';'
            else:
                return 'if(s.size()==' + str(len(key)) + '&&' + '&&'.join(conditions) + ')' + \
                    'return ' + str(value) + ';else return 0;'
        cases = {}
        for key, value in obj.items():
            ch = key.encode('utf-8')[pos]
            if ch in cases:
                cases[ch][key] = value
            else:
                cases[ch] = {key: value}
        ret = 'if(s.size()<=' + str(pos) + ')return 0;'
        ret += 'switch(s[' + str(pos) + ']){'
        for ch, o in cases.items():
            ret += 'case ' + ch2label(ch) + ':'
            ret += tostr(o, pos + 1)
        ret += 'default:return 0;}'
        return ret

    model_str = ''
    for key, value in model.items():
        if key == 'BIAS':
            model_str += 'int BIAS() const {return ' + str(value) + ';}\n'
        else:
            model_str += 'int ' + key + '(const std::string &s) const { ' + tostr(value) + '}\n'
    template = u(open('templates/cpp.hpp').read())
    template = template.replace(u'__MODEL__', model_str)
    return template

@language('java', 'java')
def java(model):
    def utf16encode(s):
        result = ''
        hexlist = '0123456789abcdef'
        s = s.encode('utf-16-be')
        try: # for python2
            l = [ord(ch) for ch in s]
        except: # for python3
            l = [ch for ch in s]

        for i in range(0, len(l), 2):
            ch1 = l[i]
            ch2 = l[i+1]
            if ch1 == 0 and 0x20 <= ch2 <= 0x7e:
                if chr(ch2) == '"':
                    result += '\\"'
                else:
                    result += chr(ch2)
            else:
                result += '\\u'
                result += hexlist[ch1>>4] + hexlist[ch1&0xF]
                result += hexlist[ch2>>4] + hexlist[ch2&0xF]
        return result

    model_str = 'public class Model {\n'
    model_str += 'public int BIAS;\n'
    model_str += 'public Map<String, Integer> '
    model_str += ','.join(key for key in model.keys() if key != 'BIAS') + ';\n'
    model_str += '}\n'

    for key, value in model.items():
        if key == 'BIAS':
            model_str += 'private static int BIAS = ' + str(value) + ';\n'
        else:
            a = list(value.items())
            a.sort()
            model_str += 'private static final String[] ' + key + '_KEYS = {' + ','.join(['"' + utf16encode(k) + '"' for k,v in a]) + '};\n'
            model_str += 'private static final int[] ' + key + '_VALUES = {' + ','.join([str(v) for k,v in a]) + '};\n'

    model_str += 'private Model newModel() {\n'
    model_str += 'Model m = new Model();\n'
    model_str += 'm.BIAS = BIAS;\n'
    for key in model.keys():
        if key == 'BIAS':
            continue
        model_str += 'm.' + key + ' = new HashMap<String, Integer>();\n'
        model_str += 'for(int i=0;i<' + key + '_KEYS.length;++i)'
        model_str += 'm.' + key + '.put(' + key + '_KEYS[i],' + key + '_VALUES[i]);\n'
    model_str += 'return m;\n'
    model_str += '}\n'

    template = open('templates/java.java').read()
    template = template.replace('__MODEL__', model_str)
    return template

@language('tex', 'sty')
def tex(model):

    def tostr(category, values):
        result = ''
        category = category.replace('1', 'i')
        category = category.replace('2', 'ii')
        category = category.replace('3', 'iii')
        category = category.replace('4', 'iv')
        category = category.replace('5', 'v')
        category = category.replace('6', 'vi')
        for key, value in values.items():
            result += r'\def\tseg@' + category + '@' + key + '{' + str(value) + '}\n'
        return result

    model_str = ''
    for key, value in model.items():
        if key == 'BIAS':
            continue
        else:
            model_str += tostr(key, value)


    template = u(open('templates/tex.sty').read())
    template = template.replace(u'__BIAS__', u"{0}".format(model['BIAS']))
    template = template.replace(u'__MODEL__', model_str)
    return template

@language('vim', 'vim')
def vim(model):
    def tostr(obj):
        items = []
        for key, value in obj.items():
            if isinstance(value, dict):
                value = tostr(value)
            key = key.replace("'", "''")
            key = "'" + key + "'"
            items.append(u'{0}:{1}'.format(key, value))
        items.sort()
        return '{' + ','.join(items) + '}'
    template = u(open('templates/vim.vim').read())
    template = template.replace(u'__MODEL__', tostr(model))
    return template

@language('go', 'go')
def go(model):
    def escape(key):
        return key.replace(u"'", u"\\'")

    def unigram2map(obj):
        items = []
        for key, value in obj.items():
            rune = u"'" + key + u"'"
            if key in [u'B1', u'B2', u'B3', u'E1', u'E2', u'E3']:
                rune = key
            elif len(key) != 1:
                continue
            items.append(u"{0}:{1}".format(rune, value))
        items.sort()
        return u'map[rune]int{' + u','.join(items) + u'}'

    def asciingram2map(obj, size):
        items = []
        for key, value in obj.items():
            keyhash = 0
            for ch in key:
                keyhash = (keyhash << 8) | ord(ch)
            items.append(u"0x{0:X}:{1}".format(keyhash, value))
        items.sort()
        return 'map[uint32]int{' + ','.join(items) + '}'

    def bigram2map(obj):
        items = []
        for key, value in obj.items():
            if len(key) != 2:
                continue
            items.append(u"0x{0:X}:{1}".format((ord(key[0]) << 32) + ord(key[1]), value))
        items.sort()
        return 'map[uint64]int{' + ','.join(items) + '}'

    def ngram2map(obj, size):
        items = []
        for key, value in obj.items():
            runes = []
            while len(key) != 0:
                if len(key) != size and len(key) >= 2 and key[0:2] in ['B1', 'B2', 'B3', 'E1', 'E2', 'E3']:
                    runes.append(key[0:2])
                    key = key[2:]
                else:
                    runes.append("'" + escape(key[0]) + "'")
                    key = key[1:]
            if len(runes) != size:
                continue
            items.append(u"[{0}]rune{{{1}}}:{2}".format(size, ','.join(runes), value))
        items.sort()
        return 'map[[' + str(size) + ']rune]int{' + ','.join(items) + '}'

    model_str = u"\n\t\tBias: {0},\n".format(model['BIAS'])
    for key in ['UP1', 'UP2', 'UP3', 'UC1', 'UC2', 'UC3', 'UC4', 'UC5', 'UC6']:
        model_str += "\t\t" + key + ": " + unigram2map(model[key]) + ",\n"
    for key in ['BP1', 'BP2', 'BC1', 'BC2', 'BC3', 'UQ1', 'UQ2', 'UQ3']:
        model_str += "\t\t" + key + ": " + asciingram2map(model[key], 2) + ",\n"
    for key in ['TC1', 'TC2', 'TC3', 'TC4', 'BQ1', 'BQ2', 'BQ3', 'BQ4']:
        model_str += "\t\t" + key + ": " + asciingram2map(model[key], 3) + ",\n"
    for key in ['TQ1', 'TQ2', 'TQ3', 'TQ4']:
        model_str += "\t\t" + key + ": " + asciingram2map(model[key], 4) + ",\n"

    for key in ['UW1', 'UW2', 'UW3', 'UW4', 'UW5', 'UW6']:
        model_str += "\t\t" + key + ": " + unigram2map(model[key]) + ",\n"
    for key in ['BW1', 'BW2', 'BW3']:
        model_str += "\t\t" + key + ": " + bigram2map(model[key]) + ",\n"
    for key in ['TW1', 'TW2', 'TW3', 'TW4']:
        model_str += "\t\t" + key + ": " + ngram2map(model[key], 3) + ",\n"

    model_str += "\t"
    template = u(open('templates/go.go').read())
    template = template.replace(u'__MODEL__', model_str)
    return template

@language('jsx', 'jsx')
def jsx(model):
    def tostr(obj):
        BIAS = ''
        items = []
        for key, value in obj.items():
            if key == 'BIAS':
                BIAS = str(value)
                continue
            if isinstance(value, dict):
                value = tostr(value)[0]
            if not re.match('[a-zA-Z0-9]+$', key):
                key = key.replace("'", "\\'")
                key = "'" + key + "'"
            items.append(u"{0}:{1}".format(key, value))
        items.sort()
        return ['{' + ','.join(items) + '}', BIAS]
    model_str, bias = tostr(model)
    template = u(open('templates/jsx.jsx').read())
    template = template.replace(u'__MODEL__', model_str).replace(u'__BIAS__', bias)
    return template

@language('csharp', 'cs')
def csharp(model):
    def tostr(obj):
        BIAS = ''
        items = []
        for key, value in obj.items():
            if key == 'BIAS':
                BIAS = str(value)
                continue

            items2 = []
            for key2, value2 in value.items():
                key2 = key2.replace("\"", "\\\"")
                key2 = "\"" + key2.replace("\"", "\\\"") + "\""
                items2.append('{' + key2 + ', ' + str(value2) + '}')
            items2.sort()

            items.append('        public static readonly Dictionary<string, int> ' + key + ' = new Dictionary<string, int>() { ' + ','.join(items2) + ' };')
        items.sort()
        return ['\n'.join(items), BIAS]
    model_str, bias = tostr(model)
    template = u(open('templates/csharp.cs').read())
    template = template.replace(u'__MODEL__', model_str).replace(u'__BIAS__', bias)
    return template

@language('julia', 'jl')
def julia(model):
    def escape(key):
        return key.replace(u'"', u'\\"')

    def unigram2dict(obj):
        items = []
        for key, value in obj.items():
            char = u"'" + key + u"'"
            if key in [u'B1', u'B2', u'B3', u'E1', u'E2', u'E3']:
                char = key
            elif len(key) != 1:
                continue
            items.append(u"{0} => {1}".format(char, value))
        items.sort()
        return u'Dict{Char,Int}(' + u', '.join(items) + u')'

    def asciingram2dict(obj, size):
        items = []
        for key, value in obj.items():
            items.append(u"t\"{0}\" => {1}".format(key, value))
        items.sort()
        return 'Dict{Tuple{' + ','.join(['UInt8'] * size) + '},Int}('+ ', '.join(items) + ')'

    def asciiunigram2dict(obj):
        items = []
        for key, value in obj.items():
            if len(key) != 1:
                continue
            items.append(u"'{0}' => {1}".format(key, value))
        items.sort()
        return 'Dict{UInt8,Int}(' + ', '.join(items) + ')'

    def ngram2dict(obj, size):
        items = []
        for key, value in obj.items():
            chars = []
            offset_flag = False
            while len(key) != 0:
                if len(key) != size and len(key) >= 2 and key[0:2] in ['B1', 'B2', 'B3', 'E1', 'E2', 'E3']:
                    chars.append(key[0:2])
                    key = key[2:]
                    offset_flag = True
                else:
                    chars.append("'" + escape(key[0]) + "'")
                    key = key[1:]
            if len(chars) != size:
                continue
            if offset_flag:
                items.append(u"({0}) => {1}".format(','.join(chars), value))
                offset_flag = False
            else:
                items.append(u"t\"{0}\" => {1}".format(''.join([c.strip("'") for c in chars]), value))
        items.sort()
        return 'Dict{Tuple{' + ','.join(['Char'] * size) + '},Int}(' + ', '.join(items) + ')'

    model_str = u"\nconst BIAS = {0}\n".format(model['BIAS'])
    for key in ['UP1', 'UP2', 'UP3', 'UC1', 'UC2', 'UC3', 'UC4', 'UC5', 'UC6']:
        model_str += "const " + key + " = " + asciiunigram2dict(model[key]) + "\n"
    for key in ['BP1', 'BP2', 'BC1', 'BC2', 'BC3', 'UQ1', 'UQ2', 'UQ3']:
        model_str += "const " + key + " = " + asciingram2dict(model[key], 2) + "\n"
    for key in ['TC1', 'TC2', 'TC3', 'TC4', 'BQ1', 'BQ2', 'BQ3', 'BQ4']:
        model_str += "const " + key + " = " + asciingram2dict(model[key], 3) + "\n"
    for key in ['TQ1', 'TQ2', 'TQ3', 'TQ4']:
        model_str += "const " + key + " = " + asciingram2dict(model[key], 4) + "\n"

    for key in ['UW1', 'UW2', 'UW3', 'UW4', 'UW5', 'UW6']:
        model_str += "const " + key + " = " + unigram2dict(model[key]) + "\n"
    for key in ['BW1', 'BW2', 'BW3']:
        model_str += "const " + key + " = " + ngram2dict(model[key], 2) + "\n"
    for key in ['TW1', 'TW2', 'TW3', 'TW4']:
        model_str += "const " + key + " = " + ngram2dict(model[key], 3) + "\n"

    template = u(open('templates/julia.jl').read())
    template = template.replace(u'__MODEL__', model_str)
    return template

def main():
    parser = OptionParser()
    parser.add_option("-s", "--scale", dest = "scale",
                      default = 0, type = "int",
                      help="The scale of parameters")
    parser.add_option("-d", "--dir", dest = "dir",
                      default = '.',
                      help="output directory")
    parser.add_option("-o", "--output", dest = "output",
                      default = '',
                      help="output filename")
    parser.add_option("-m", "--model", dest = "model",
                      default = '',
                      help="model filename")
    (options, args) = parser.parse_args()
    if options.model:
        model = load(open(options.model))
    else:
        model = load(sys.stdin)
    lang = args[0]
    if lang == 'all':
        for maker in languages.values():
            with open(options.dir + '/tinysegmenter.' + maker.ext, 'w') as f:
                f.write(fileformat(maker.execute(model)))
    elif lang in languages:
        maker = languages[lang]
        filename = options.output or options.dir + '/tinysegmenter.' + maker.ext
        with open(filename, 'w') as f:
            f.write(fileformat(maker.execute(model)))
    else:
        print >>sys.stderr, "Unknow language:", lang

if __name__ == "__main__":
    main()