# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
from sklearn.model_selection import train_test_split
import math
import re
def PreData():
    txt = open(r'./BME标注.txt', encoding='utf-8').readlines()
    # 设置random_state保证每次划分都一样
    train_data_, test_data_ = train_test_split(txt, test_size=0.2, random_state=10)
    # print(train_data[:10])
    train_data = []
    test_data = []
    for iter in train_data_:
        tmp =  re.split('。|？|！',iter)
        train_data += tmp


    for iter in test_data_:
        iter = re.sub(' *','',iter)
        iter = re.sub('/[A-Z]*','',iter)
        tmp =  re.split('。|？|！',iter)
        test_data += tmp

    return train_data,test_data
        
def SeqData(train_data):
    seq_train = []
    # seq_test = []
    for iter in train_data:
        tmp =  re.split('。|？|！',iter)
        st = ''
        for ech in tmp:
            seq = re.findall('/[A-Z]',ech)
            st += ''.join(seq)
        seq_train.append(st)

    return seq_train



# +
# test_0,test_1 = PreData()
# # test_2 = SeqData(test_0,)
# test_0[:5]
# s = 'abc/dfe/ee'
# print (s.split('/'))
# -

def Count(train_data_str):
    label_dic = {}
    status = ['B','M','S','E']
    for iter in status:
        label_dic[iter] = 0
    for iter in range(len(train_data_str)):
        label_dic[train_data_str[iter]] += 1
    
    B = label_dic['B']
    M = label_dic['M']
    S = label_dic['S']
    E = label_dic['E']
    return B,M,S,E


# +

def start_pro(train_data):
    start_p = {}
    status = ['B','M','S','E']
    for iter in status:
        start_p[iter] = 0
    for iter in train_data:
        # print (type(iter))
        iter = re.sub('/','',iter)
        # 初始概率
        if len(iter) == 0:
            continue
        word_b = iter[0]
        if word_b in start_p:
            start_p[word_b] += 1
    
    for iter in start_p.keys():
        if start_p[iter] < 1e-7:
            start_p[iter] += 1
        start_p[iter] = -math.log(start_p[iter]/len(train_data))
    return start_p


# -

def transfer_pro(train_data):
    train_data_s = ''.join(train_data)
    train_data_s = re.sub('/s','',train_data_s)
    train_data_s = re.sub('/','',train_data_s)
    transfer_p = {}        

    status = ['BM','MB','MS','SM','ME','EM','BS','SB','BE','EB','ES','SE','BB','MM','SS','EE']
    for iter in status:
        transfer_p[iter] = 0
    
    for iter in range(len(train_data_s)):
        if iter == 0:
            continue
        transfer_p[train_data_s[iter-1:iter +1]] += 1
    for iter in transfer_p.keys():
        if transfer_p[iter]  < 1e-7:
            transfer_p[iter] += 1
        transfer_p[iter] = -math.log(transfer_p[iter]/len(train_data_s))
        
    return transfer_p

# +
# s = "本/B报/E  布/B鲁/M塞/M尔/E "
# List = re.findall('./E',s)
# List[0] = re.sub('/E','',List[0])
# print (List[0])

# +


def knowndic(train_data):
    knowledge_table = {}
    for sentence in train_data:
        tmp = re.findall('./[A-Z]',sentence)
        for iter in tmp:
            if iter[0] in knowledge_table:
                if iter[-1] not in knowledge_table[iter[0]]:
                    knowledge_table[iter[0]].append(iter[-1])
                    # print(type(knowledge_table[iter[0]]))
            else:
                knowledge_table[iter[0]] = [iter[-1]]
                # knowledge_table[iter[0]] = {iter[-1]:1}
    return knowledge_table



# -

s = ['我/B的/E我/S们/M的/E你/B和/M我/B']
knowndic(s)


def _launch_pro(train_data,tar):
    dit = {}
    p_tar = '.' + tar
    _list = re.findall(p_tar,train_data)
    for iter in _list:
        iter = re.sub(tar,'',iter)
        if iter in dit:
            dit[iter] += 1
        else:
            dit[iter] = 1
    
    for k in dit.keys():
        if dit[k] <  1e-7:
            dit[k] += 1
        dit[k] = -math.log(dit[k] / len(_list))
    return dit


def launch_pro(train_data):
    train_data = ''.join(train_data)
    launch_p = {
        'B':{},
        'M':{},
        'S':{},
        'E':{}
    }
    for iter in launch_p.keys():
        launch_p[iter] = _launch_pro(train_data,'/'+iter)
    
    return launch_p


s = "本/B报/E  布/B鲁/M塞/M尔/E \n"
p = launch_test = launch_pro(s)
print(p)

# +
# 输出概率字典到json文件
import json
def WriteF(path,dic):
    with open(path,'w',encoding='utf-8') as f:
        # ensuer_ascii Fales 不用编码输出，显示为汉字
        f.write(json.dumps(dic,sort_keys=True,ensure_ascii=False,indent=4, separators=(',', ':')))
    return

def WriteDataF(path,data):
    with open(path,'w',encoding='utf-8') as f:
        for iter in data:
            f.writelines(iter)
        f.close()
    return


# +
# viterbi算法的实现
# https://www.zhihu.com/question/20136144

def viterbi(start_p, trans_p, emit_p,states,obs):  # 维特比算法（一种递归算法）
    V = [{}]
    path = {}
    for y in states:
        V[0][y] = start_p[y] * emit_p[y].get(obs[0], 0)  # 在位置0，以y状态为末尾的状态序列的最大概率
        path[y] = [y]

    for t in range(1, len(obs)):
        V.append({})
        newpath = {}
        for y in states:
            state_path = ([(V[t - 1][y0] * trans_p[y0].get(y, 0) * emit_p[y].get(obs[t], 0), y0) for y0 in states if V[t - 1][y0] > 0])
            if state_path == []:
                (prob, state) = (0.0, 'S')
            else:
                (prob, state) = max(state_path)
            V[t][y] = prob
            newpath[y] = path[state] + [y]

        path = newpath  # 记录状态序列
    (prob, state) = max([(V[len(obs) - 1][y], y) for y in states])  # 在最后一个位置，以y状态为末尾的状态序列的最大概率
    return (prob, path[state])  # 返回概率和状态序列

# 分词
def cut(sent, prob_start, prob_trans, prob_emit):
    prob, pos_list = viterbi(sent, ('B', 'M', 'E', 'S'), prob_start, prob_trans, prob_emit)
    seglist = list()
    word = list()
    for index in range(len(pos_list)):
        if pos_list[index] == 'S':
            word.append(sent[index])
            seglist.append(word)
            word = []
        elif pos_list[index] in ['B', 'M']:
            word.append(sent[index])
        elif pos_list[index] == 'E':
            word.append(sent[index])
            seglist.append(word)
            word = []
    seglist = [''.join(tmp) for tmp in seglist]

    return seglist



# +
def read_j(path):
    dic = {}
    with open(path, 'r') as f:
        dic = json.load(fp=f)
    return dic

def test_viterbi():
    start_p = read_j('./初始概率.json')
    transfer_p = read_j('./转移概率.json')
    launch_p = read_j('./发射概率.json')
    knowndic = read_j('./知识表.json')
    s = '抓典型是我们党的优良传统'
    print(viterbi(start_p,transfer_p,launch_p,knowndic,s))
    
    

# test_viterbi()

# +

def main():
    train_data,test_data = PreData()
    WriteDataF('./训练集.txt',train_data)
    WriteDataF('./测试集.txt',test_data)
    seq_txt = SeqData(train_data)
    know = knowndic(train_data)
    start_p = start_pro(seq_txt)
    transfer_p = transfer_pro(seq_txt)
    launch_p = launch_pro(train_data)
    WriteF('./知识表.json',know)
    WriteF('./初始概率.json',start_p)
    WriteF('./转移概率.json',transfer_p)
    WriteF('./发射概率.json',launch_p)

    s = '如何选择联网漫游台'
    print(viterbi(know,transfer_p,launch_p,start_p,s))
if __name__ == '__main__':
    main()
# -

# # 老师，我处理的知识表里面有一些字符，viterbi算法进行不下去

StarcoderdataJupyterScriptsDedupFiltered