andreslu/orion

1from ngram import NGram2 3 4def post_process_template(tB):5    if tB.endswith('.') == False:6        tB += '.'7    return tB8    # return tB.split('.')[0] + '.'9 10 11def construct_template(words, templateA, if_then=False):12    if len(words) >= 2:13        templates = ['{} <mask> '.format(words[0])]14        for i in range(1, len(words)-1):15            templates[0] += '{} <mask> '.format(words[i])16        templates[0] += '{}.'.format(words[-1])17    elif len(words) == 1:18        templates = [19            # '{} is <mask>.'.format(words[0]),20            '{} <mask>.'.format(words[0])]21 22    elif len(words) == 0:23        templates = []24 25    if if_then:26        for word in words:27            index = templateA.index('<mask>')28            templateA = templateA[:index] + word + templateA[index + len('<mask>'):]29        templates = ['If ' + templateA + ' then ' + template for template in templates]30 31    return templates32 33 34def filter_words(words_prob):35    word_count = {}36    token1_count = {}37    word2_count = {}38    ret = []39    for words, prob, *_ in words_prob:40        filter_this = False41 42        # filter repetitive token43        token_count = {}44        for word in words:45            for token in word.split(' '):46                if token in token_count:47                    filter_this = True48                token_count[token] = 149        if filter_this:50            prob *= 0.551 52        # filter repetitive words53        if len(words) == 2 and words[0] == words[1]:54            continue55 56        # filter repetitive first token57        token1 = words[0].split(' ')[0]58        if token1 not in token1_count:59            token1_count[token1] = 160        else:61            token1_count[token1] += 162            prob /= token1_count[token1]63 64        for word in words:65            if word not in word_count:66                word_count[word] = 067            word_count[word] += 168            prob /= word_count[word]69 70        if len(words) == 2:71            if words[1] not in word2_count:72                word2_count[words[1]] = 073            word2_count[words[1]] += 174            prob /= word2_count[words[1]]75 76        ret.append([words, prob])77    return sorted(ret, key=lambda x: x[1], reverse=True)78 79 80import math81from copy import deepcopy82 83 84def convert_for_print(arr):85    ret = deepcopy(arr)86    for i in range(len(ret)):87        ret[i][1] = round(ret[i][1], 7)88        if len(ret[i]) == 3:89            for j in range(len(ret[i][2])):90                ret[i][2][j] = round(ret[i][2][j], 7)91    return ret92 93 94def formalize_tA(tA):95    tA = tA.strip()96    if tA.endswith('.'):97        tA = tA[:-1].strip() + '.'98    else:99        tA += '.'100    tA = tA.replace(' ,', ',')101    tA = tA.replace(" '", "'")102    return tA103 104 105ngram_n = 3106 107 108def extract_similar_words(txt, words):109    max_word_length = 0110    for word in words:111        if len(word) > max_word_length:112            max_word_length = len(word)113 114    txt_ngrams = []115    for i in range(len(txt)):116        for j in range(i + ngram_n, min(len(txt), i + max_word_length + 5)):117            txt_ngrams.append(txt[i:j].lower())118    n = NGram(txt_ngrams, key=lambda x: x.lower(), N=ngram_n)119    ret = []120    for word in words:121        matched_word = n.find(word.lower(), 0.5)122        if matched_word is None:123            return None124        ret.append(matched_word)125    return ret126 127 128def extract_words(txt, words):129    for word in words:130        if word not in txt:131            return None132    return [word.lower() for word in words]133