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

# # Perplexidade
# 
# Uma das formas de avaliar um modelo de linguagem é calculando a sua perplexidade.
# 
# Dado um conjunto de testes com $m$ frases ($s_1, s_2, ..., s_m$), um modelo é bom se ele atribui probabilidade alta às frases do conjunto de testes, em outras palavras, se ele consegue prever as frases.
# 
# Vamos utilizar um dataset com os romances de Machado de Assis. O código abaixo lê o arquivo e extrai o conteúdo textual das obras, excluindo os nomes dos capítulos.

# In[1]:


filepath = "machadodeassiscorpus.txt"
text = ""

with open(filepath) as corpus:
    for line in corpus:
        if not (line.startswith("#") or line.startswith("%")):
            text += line
            
text = text.lower()


# Os primeiros 100 caracteres do dataset são:

# In[2]:


print(text[:100])


# As probabilidades são calculadas usando tokens do texto. Para quebrar o texto em tokens utilizaremos a biblioteca NLTK.

# In[3]:


from nltk.tokenize import word_tokenize

tokens = word_tokenize(text)


# Os 50 primeiros tokens são:

# In[4]:


print(tokens[:50])


# Precisamos agrupar os tokens em frases para calcular as probabilidades. Nesse caso, consideraremos o fim de uma frase quando o token for ".", "?" ou "!".

# In[5]:


sentences = []
temp_sentence = []

for token in tokens:
    temp_sentence.append(token)

    if token in ".!?":
        sentences.append(temp_sentence)
        temp_sentence = []


# Vejamos as duas primeiras frases retornadas:

# In[6]:


print(" ".join(sentences[0]))
print(" ".join(sentences[1]))


# Para calcular a probabilidade que o modelo atribui a todo conjunto de testes multiplicaremos as probabilidades atribuídas a cada frase:
# 
# \begin{equation*}
# P(teste) = \prod_{i=1}^{m}p(s_i)
# \end{equation*}
# 
# O valor calculado rapidamente se torna muito pequeno, podendo causar problemas de underflow. Por isso, usaremos o $\log_2$ da probabilidade. Como $\log (a*b) = \log a + \log b$, temos:
# 
# 
# 
# \begin{equation*}
# \log_2 P(teste) = \sum_{i=1}^{m}\log_2 p(s_i)
# \end{equation*}
# 

# Primeiramente vamos utilizar unigramas. Com isso, a probabilidade de um token $w_i$ é calculada assim:
# 
# \begin{equation*}
# P(w_i | w_1 w_2 ... w_{i-1}) = P(w_i)
# \end{equation*}
# 
# Dessa forma, precisamos contar a ocorrência de cada token no dataset.

# In[7]:


unigrams_count = {}

for token in tokens:
    unigrams_count[token] = unigrams_count.get(token, 0) + 1


# Podemos ver no resultado da contagem, por exemplo, quantas vezes a palavra "dia" e "mulher" aparecem.

# In[8]:


print("Contagem da palavra 'dia': {}".format(unigrams_count["dia"]))
print("Contagem da palavra 'mulher': {}".format(unigrams_count["mulher"]))


# A probabilidade de uma frase é a multiplicação das probabilidades dos seus tokens. Nesse caso, se a frase tem $n$ tokens, temos:
# 
# \begin{equation*}
# P(s) = P(w_1) * P(w_2) * ... * P(w_n)
# \end{equation*}

# Entretanto, dependendo da frase, esse valor já pode sofrer com problemas de underflow. Por isso, vamos usar o $\log_2$ nesses cálculos.
# 
# \begin{equation*}
# \log_2 P(s) = \log_2 P(w_1) + \log_2 P(w_2) + ... + \log_2 P(w_n)
# \end{equation*}

# Para calcular a probabilidade de uma frase definiremos uma função. Ela vai receber a frase, os tokens e as contagens como parâmetros. 
# 
# As contagens podem ser computadas através dos tokens, mas, por questões de desempenho, vamos passar para função as contagens que já temos.

# In[9]:


from math import log2

def sentence_unigram_prob(sentence, tokens, counts):
    total = len(tokens)

    prob = 0
    for token in sentence:
        token_prob = counts[token]/total
        prob += log2(token_prob)

    return prob


# Testando com frases do dataset:

# In[10]:


print("Log_2 da probabilidade da frase '{}': {}".format(" ".join(sentences[0]),
    sentence_unigram_prob(sentences[0], tokens, unigrams_count)))

print("Log_2 da probabilidade da frase '{}': {}".format(" ".join(sentences[1]),
    sentence_unigram_prob(sentences[1], tokens, unigrams_count)))


# Agora vamos calcular o $\log_2$ da probabilidade que o modelo atribui ao conjunto de testes:

# In[11]:


unigram_prob = 0
for sentence in sentences:
    sentence_prob = sentence_unigram_prob(sentence, tokens, unigrams_count)
    unigram_prob += sentence_prob
    
print(unigram_prob)


# Um dataset com muitas frases tende a possuir uma probabilidade menor que um com poucas, pois mais valores de probabilidade são multiplicados. Por isso, para normalizar o resultado, vamos tirar uma média dividindo o resultado pelo número de tokens.

# In[12]:


total_tokens = len(tokens)
mean = unigram_prob/total_tokens
print(mean)


# Para finalizar, vamos elevar 2 ao valor negativo da média. Assim, a perplexidade será um número positivo e quanto menor ela for, maior é a probabilidade atribuída pelo modelo.
# 
# \begin{equation*}
# Perplexidade = 2^{-\frac{1}{M}\sum_{i=1}^{m}\log_2 p(s_i)}
# \end{equation*}
# 
# 

# In[13]:


perplexity = 2**(-1*(mean))
print(perplexity)


# O modelo unigrama tende a ser ruim em prever probabilidades, pois cada token é avaliado individualmente, sem considerar tokens anteriores. Espera-se que um modelo bigrama ou trigrama tenham desempenho melhores respectivamente. Assim, vamos calcular a perplexidade usando bigramas.

# O primeiro passo é fazer a contagem dos bigramas:

# In[14]:


bigrams_count = {}

for i in range(1, len(tokens)):
    pair = (tokens[i-1], tokens[i])
    bigrams_count[pair] = bigrams_count.get(pair, 0) + 1


# Podemos ver no resultado da contagem, por exemplo, quantas vezes os pares ("o", "dia") e ("a", "mulher") aparecem

# In[15]:


print("Contagem do par ('o', 'dia'): {}".format(bigrams_count[("o", "dia")]))
print("Contagem da par ('a', 'mulher'): {}".format(bigrams_count[("a", "mulher")]))


# Utilizando bigramas, a probabilidade de um token $w_i$ é calculada assim:
# 
# \begin{equation*}
# P(w_i | w_1 w_2 ... w_{i-1}) = P(w_i | w_{i-1})
# \end{equation*}
# 
# Com isso, a probabilidade de uma frase com $n$ tokens é:
# 
# \begin{equation*}
# P(s) = P(w_1) * P(w_2 | w_1) * P(w_3 | w_2) ... * P(w_n | w_{n-1})
# \end{equation*}
# 
# Sendo
# 
# \begin{equation*}
# P(b | a) = \frac{count(a, b)}{count(b)}
# \end{equation*}
# 
# Calcularemos o $\log_2$ das probabilidades usando a função a seguir:

# In[16]:


def sentence_bigram_prob(sentence, tokens, unigram_counts, bigram_counts):
    total = len(tokens)
    
    prob = log2(unigram_counts[tokens[0]]/total)

    for i in range(1, len(sentence)):
        pair = (tokens[i-1], tokens[i])
        token_prob = bigram_counts[pair]/unigram_counts[tokens[i-1]]

        prob += log2(token_prob)

    return prob


# Testando com o dataset:

# In[17]:


print("Log_2 da probabilidade da frase '{}': {}".format(" ".join(sentences[0]),
    sentence_bigram_prob(sentences[0], tokens, unigrams_count, bigrams_count)))

print("Log_2 da probabilidade da frase '{}': {}".format(" ".join(sentences[1]),
    sentence_bigram_prob(sentences[1], tokens, unigrams_count, bigrams_count)))


# Agora vamos calcular o $\log_2$ da probabilidade que o modelo bigrama atribui ao conjunto de testes:

# In[18]:


bigram_prob = 0
for sentence in sentences:
    sentence_prob = sentence_bigram_prob(sentence, tokens, unigrams_count, bigrams_count)
    bigram_prob += sentence_prob
    
print(bigram_prob)


# Finalizando o cálculo da perplexidade, temos:

# In[19]:


mean = bigram_prob/total_tokens
perplexity = 2**(-1*(mean))
print(perplexity)


# Como esperado, a perplexidade para bigramas é menor, mostrando que esse modelo é melhor que o unigrama.