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

# <p>
# <img src="../imgs/EII-ULPGC-logo.jpeg" width="430px" align="right">

# # **NOTEBOOK 20**
# ---

# 
# 
# # **Modelos del lenguaje basados en redes neuronales artificiales**

# ## **BERT (Bidirectional Encoder Representations from Transformers)**
# 
# "*BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding*"
# https://arxiv.org/pdf/1810.04805.pdf
# 
# 
# Desarrollado por Google en 2018, marcó un hito significativo en el campo del NLP. Este modelo se basa en la arquitectura Transformer, introducida en 2017 por Vaswani et al., y es conocido por su capacidad de entender el contexto de las palabras en el texto de una manera más sofisticada que los modelos anteriores. A diferencia de los enfoques anteriores que procesaban el texto de manera secuencial (de izquierda a derecha o viceversa), BERT analiza el texto en ambas direcciones simultáneamente. Esta bidireccionalidad permite a BERT capturar el contexto completo de una palabra, mirando tanto el texto anterior como el siguiente. Esto lo hace particularmente eficaz para entender el significado y la intención detrás de las palabras en oraciones complejas.
# 
# BERT se preentrena en un corpus masivo de texto no etiquetado y utiliza dos estrategias principales:
# 
# - **Masked Language Model (MLM):** donde se ocultan aleatoriamente palabras del texto y el modelo aprende a predecirlas.
# 
# - **Next Sentence Prediction (NSP):** donde el modelo aprende a predecir si una oración es la continuación lógica de otra.
# 
# <img src="./imgs/BERT_MODEL.svg">
# 
# ### **1. Masked Language Model (MLM)**
# 
# 1. **Enmascaramiento de Tokens**:
#    - Alrededor del 15% de los tokens en cada secuencia de entrada se seleccionan al azar para ser enmascarados.
#    - Estos tokens seleccionados se reemplazan con un token especial `[MASK]`.
# 
# 2. **Diversificación en el Enmascaramiento**:
#    - No todos los tokens seleccionados se enmascaran de la misma manera:
#      - Un 80% de las veces, el token seleccionado se reemplaza realmente por `[MASK]`.
#      - Un 10% de las veces, se reemplaza con un token aleatorio.
#      - En el 10% restante, el token se deja sin cambios.
# 
#      Al diversificar el enmascaramiento, se asegura que el modelo no se sobreajuste a los tokens [MASK] y que aprenda a utilizar el contexto para predecir palabras, lo cual es más representativo de cómo se utilizará el modelo en aplicaciones del mundo real.
# 
#      Al dejar algunos tokens seleccionados sin cambios, BERT aprende a predecir la idoneidad de las palabras actuales en su contexto original.
# 
#      La sustitución con tokens aleatorios ayuda a que el modelo sea robusto frente a entradas inesperadas o ruidosas, mejorando su capacidad para manejar errores o variaciones en los datos de entrada.
# 
# 3. **Predicción de Tokens Enmascarados**:
#    - El objetivo del modelo durante el entrenamiento es predecir los tokens originales de aquellos que han sido enmascarados o alterados.
#    - Esto enseña a BERT a entender el contexto y la relación entre las palabras en una secuencia.
# 
# Durante el entrenamiento de BERT, solo los tokens seleccionados para predicción (aproximadamente el 15% de los tokens en cada secuencia de entrada) contribuyen al cálculo de la función de pérdida. Esto significa que el modelo únicamente calcula el error para estos tokens específicos, permitiéndole aprender a "adivinar" palabras en función de su contexto sin procesar cada token de la secuencia.
# 
# Dentro de estos tokens seleccionados, algunos se reemplazan por el token [MASK], otros se sustituyen por una palabra aleatoria, y un pequeño porcentaje se deja sin cambios. Aunque estos tokens no se alteran, la diferencia clave es que el modelo sí intenta predecirlos como parte de los tokens seleccionados para la predicción. Esto contrasta con los tokens no seleccionados (aquellos que no se modifican ni se predicen), los cuales no participan en el cálculo de la pérdida y, por lo tanto, el modelo no tiene que aprender nada sobre ellos.
# 
# Este enfoque, donde algunos tokens se dejan "sin cambios" pero aún participan en el entrenamiento, enseña al modelo a inferir palabras en su contexto sin depender exclusivamente de "pistas" explícitas de enmascaramiento, lo que mejora su capacidad para manejar el lenguaje en situaciones del mundo real.
# 
# ### **2. Next Sentence Prediction (NSP)**
# 
# 1. **Entendimiento de Relaciones entre Oraciones**:
#    - NSP es una tarea de clasificación binaria para predecir si una oración B es la continuación lógica de una oración A.
# 
# 2. **Preparación de Datos de Entrenamiento**:
#    - Durante el preentrenamiento, se toman pares de oraciones de un corpus de texto.
#    - En aproximadamente el 50% de los casos, la oración B es de hecho la oración que sigue naturalmente a la oración A.
#    - En el otro 50%, la oración B es una oración aleatoria del corpus, no relacionada con la oración A.
# 
# 3. **Uso del Token `[CLS]`**:
#    - Para cada par de oraciones, el token `[CLS]` se añade al principio, y un token `[SEP]` se utiliza para separar las dos oraciones.
#    - El modelo utiliza la representación del token `[CLS]` para hacer la predicción.
# 
# ### Importancia del Preentrenamiento
# 
# 1. **Aprendizaje de Contexto y Relaciones**:
#    - MLM y NSP juntos enseñan a BERT a entender tanto el contexto a nivel de palabra (a través de MLM) como las relaciones entre oraciones completas (a través de NSP).
# 
# 
# 2. **Aplicabilidad en Tareas de NLP**:
#    - Una vez preentrenado, BERT puede ser afinado con un conjunto de datos más pequeño para tareas específicas de NLP, como la clasificación de texto, la respuesta a preguntas y el reconocimiento de entidades nombradas.
# 

# ### **BERT - 🤗 Hugging Face: Classificación de textos**
# 
# Vamos a crear un modelo BERT con la librería HuggingFace que nos permita utilizar modelos preentrenados y re-entrenarlos para las tareas que queramos. En este caso, vamos a realizar la clasificación de noticias en tres categorías: deportes, cultura y política.

# #### **Modelo**

# In[1]:


from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained('bert-base-multilingual-cased')

model = AutoModelForSequenceClassification.from_pretrained("bert-base-multilingual-cased", num_labels=3)


# Revisa el modelo que acabamos de crear. Comprueba que los elementos que lo componen y su estructura corresponden con lo que hemos visto en la teoría. ¿A qué corresponde el *pooler*?

# In[2]:


model


# #### **Tokenizador**
# 
# Fíjate en el parámetro *add_special_tokens=False*. ¿Qué ocurre si lo cambiamos a *True*? ¿Qué son los tokens especiales? ¿Qué son los *token_type_ids*? ¿Y los *attention_mask*?

# In[3]:


result = tokenizer.encode_plus("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt")

print(result.keys())
print(result['input_ids'])
print(result['token_type_ids'])
print(result['attention_mask'])


# In[4]:


result = tokenizer.encode_plus("Hello, my dog is cute", add_special_tokens=True, return_tensors="pt")
print(result.keys())
print(result['input_ids'])
print(result['token_type_ids'])
print(result['attention_mask'])


# In[5]:


result = tokenizer.decode([101, 31178,   117, 15127, 17835, 10124, 21610, 10112, 102])
print(result)


# In[6]:


import torch
from torch.utils.data import Dataset

class TextClassificationDataset(Dataset):
    def __init__(self, filename, tokenizer, max_length=256):
        self.tokenizer = tokenizer
        self.sentences = []
        self.labels = []
        self.max_length = max_length

        with open(filename, 'r', encoding='utf-8') as file:
            for line in file:
                start = line.find('"')
                end = line.find('"', start + 1)
                sentence = line[start + 1:end].strip()
                label = int(line[end + 1:].strip()[-1])

                # Crear los tokens para alimentar a BERT de HuggingFace
                tokens = self.tokenizer.encode_plus(sentence, add_special_tokens=True, max_length=self.max_length, truncation=True, padding='max_length', return_tensors='pt')
                self.sentences.append(tokens)
                self.labels.append(label)

    def __len__(self):
        return len(self.sentences)

    def __getitem__(self, idx):
        item = {key: val for key, val in self.sentences[idx].items()}
        item['labels'] = torch.tensor(self.labels[idx])
        item['input_ids'] = item['input_ids'].squeeze()
        item['attention_mask'] = item['attention_mask'].squeeze()
        item['token_type_ids'] = item['token_type_ids'].squeeze()
        
        return item
            

# Creamos el conjunto de datos
dataset = TextClassificationDataset('data/dataset_clas_texto.txt', tokenizer)

# Separamos el conjunto de datos en entrenamiento y validación
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])


# #### **Entrenamiento**
# 
# Usaremos el objeto *Trainer* de la librería *transformers* para entrenar nuestro modelo. Para ello, necesitamos el modelo, el tokenizador, los datos de entrenamiento y los datos de validación. Todos los parámetros que necesitamos para el entrenamiento están definidos en el objeto *TrainingArguments*.

# In[7]:


from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="test_trainer",
    num_train_epochs=3,
    logging_steps=1,
    logging_strategy='steps'
)

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset
)


# In[8]:


trainer.train()


# In[9]:


import os
import matplotlib.pyplot as plt
import json

def plot_training_loss(log_dir):
    history = trainer.state.log_history
    training_loss = [log['loss'] for log in history if 'loss' in log]

    plt.plot(training_loss, label='Training Loss')
    plt.xlabel('Steps')
    plt.ylabel('Loss')
    plt.title('Training Loss Over Time')
    plt.legend()
    plt.show()

plot_training_loss('./logs')


# #### **Inferencia**
# 
# Vamos a probar nuestro modelo con algunas noticias de ejemplo.

# In[10]:


sentences = []
labels = []

with open('data/dataset_clas_texto_test.txt', 'r', encoding='utf-8') as file:
            for line in file:
                start = line.find('"')
                end = line.find('"', start + 1)
                sentence = line[start + 1:end].strip()
                label = int(line[end + 1:].strip()[-1])
                sentences.append(sentence)
                labels.append(label)                


# In[11]:


clases = {
    0: 'Deportes',
    1: 'Cultura',
    2: 'Política'
}

device = torch.device("cpu")
model = model.to(device)


# In[12]:


aciertos = 0
for sentence, label in zip(sentences, labels):
    tokens = tokenizer.encode_plus(sentence, add_special_tokens=True, max_length=256, truncation=True, padding='max_length', return_tensors='pt')
    tokens.to(device)
    outputs = model(**tokens)
    y_pred = torch.argmax(outputs.logits).item()
    y = label
    if y_pred == y:
        aciertos += 1
    print(f"{sentence} -> {clases[y]}, {clases[y_pred]}")

print("-"*50)  
print(f"Aciertos: {aciertos}, Total: {len(sentences)}, Accuracy: {aciertos / len(sentences)}")