import re, collections

  

def get_vocab(filename):

    vocab = collections.defaultdict(int)

    with open(filename, 'r', encoding='utf-8') as fhand:

        for line in fhand:

            words = line.strip().split()

            for word in words:

                vocab[' '.join(list(word)) + ' </w>'] += 1

    return vocab

  

def get_stats(vocab): # 构建字符对频数字典

    pairs = collections.defaultdict(int)

    for word, freq in vocab.items():

        symbols = word.split()

        for i in range(len(symbols)-1):

            pairs[symbols[i],symbols[i+1]] += freq

    return pairs

  

def merge_vocab(pair, v_in): # 将频率最大的字符对替换

    v_out = {}

    bigram = re.escape(' '.join(pair)) # 不转义

    p = re.compile(r'(?<!\S)' + bigram + r'(?!\S)') # 意思是bigram前面没有非空格字符，后面也没有非空格字符

    for word in v_in:

        w_out = p.sub(''.join(pair), word)

        v_out[w_out] = v_in[word]

    return v_out

  

def get_tokens(vocab):

    tokens = collections.defaultdict(int)

    for word, freq in vocab.items():

        word_tokens = word.split()

        for token in word_tokens:

            tokens[token] += freq

    return tokens

  

vocab = {'l o w </w>': 5, 'l o w e r </w>': 2, 'n e w e s t </w>': 6, 'w i d e s t </w>': 3}

print('==========')

print('Tokens Before BPE')

tokens = get_tokens(vocab)

print('Tokens: {}'.format(tokens))

print('Number of tokens: {}'.format(len(tokens)))

print('==========')

  

num_merges = 5

for i in range(num_merges):

    pairs = get_stats(vocab)

    if not pairs:

        break

    best = max(pairs, key=pairs.get)

    vocab = merge_vocab(best, vocab)

    print('Iter: {}'.format(i))

    print('Best pair: {}'.format(best))

    print(vocab)

    # tokens = get_tokens(vocab)

    # print('Tokens: {}'.format(tokens))

class WordpieceTokenizer(object):

"""Runs WordPiece tokenziation."""

    def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200):
    
        self.vocab = vocab
        
        self.unk_token = unk_token
        
        self.max_input_chars_per_word = max_input_chars_per_word
        
    def tokenize(self, text):
        
        """Tokenizes a piece of text into its word pieces.
        
        This uses a greedy longest-match-first algorithm to perform tokenization
        
        using the given vocabulary.
        
        For example:
        
        input = "unaffable"
        
        output = ["un", "##aff", "##able"]
        
        Args:
        
        text: A single token or whitespace separated tokens. This should have
        
        already been passed through `BasicTokenizer.
        
        Returns:
        
        A list of wordpiece tokens.
        
        """
        
        text = convert_to_unicode(text)
        
        output_tokens = []
        
        for token in whitespace_tokenize(text):
        
        chars = list(token)
        
        if len(chars) > self.max_input_chars_per_word:
        
        output_tokens.append(self.unk_token)
        
        continue
        
        is_bad = False
        
        start = 0
        
        sub_tokens = []
        
        while start < len(chars):
        
        end = len(chars)
        
        cur_substr = None
        
        while start < end:
        
        substr = "".join(chars[start:end])
        
        if start > 0:
        
        substr = "##" + substr
        
        if substr in self.vocab:
        
        cur_substr = substr
        
        break
        
        end -= 1
        
        if cur_substr is None:
        
        is_bad = True
        
        break
        
        sub_tokens.append(cur_substr)
        
        start = end
        
        if is_bad:
        
        output_tokens.append(self.unk_token)
        
        else:
        
        output_tokens.extend(sub_tokens)
        
        return output_tokens

def encode_word(word, model):

    # word初步分词，model中为-log值

    best_segmentations = [{"start": 0, "score": 1}] + [

        {"start": None, "score": None} for _ in range(len(word))

    ]

    for start_idx in range(len(word)):

        # This should be properly filled by the previous steps of the loop

        best_score_at_start = best_segmentations[start_idx]["score"]

        for end_idx in range(start_idx + 1, len(word) + 1):

            token = word[start_idx:end_idx]

            if token in model and best_score_at_start is not None:

                score = model[token] + best_score_at_start # 加上start，即前缀对应的损失，log中相加等于原来概率相乘

                # If we have found a better segmentation ending at end_idx, we update

                if (

                    best_segmentations[end_idx]["score"] is None

                    or best_segmentations[end_idx]["score"] > score # score即loss。越小越好

                ):

                    best_segmentations[end_idx] = {"start": start_idx, "score": score}

                print(token)

                print(best_segmentations[end_idx]["score"], score)

                print(best_segmentations)

    segmentation = best_segmentations[-1]

    if segmentation["score"] is None:

        # We did not find a tokenization of the word -> unknown

        return ["<unk>"], None

    # 从后向前的最佳路径，即维特比算法

    score = segmentation["score"]

    start = segmentation["start"]

    end = len(word)

    tokens = []

    while start != 0:

        tokens.insert(0, word[start:end])

        next_start = best_segmentations[start]["start"]

        end = start

        start = next_start

    tokens.insert(0, word[start:end])

    return tokens, score

def compute_loss(model):
    loss = 0
    for word, freq in word_freqs.items():
        _, word_loss = encode_word(word, model)
        loss += freq * word_loss
    return loss

import copy


def compute_scores(model):
    scores = {}
    model_loss = compute_loss(model)
    for token, score in model.items():
        # We always keep tokens of length 1
        if len(token) == 1:
            continue
        model_without_token = copy.deepcopy(model)
        _ = model_without_token.pop(token)
        scores[token] = compute_loss(model_without_token) - model_loss
    return scores