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

# # Déboguer le pipeline d'entraînement
# 
# Ce chapitre portant sur le débogage, la langue nous importe peu ici. Nous nous intéressons surtout à la logique du code pour comprendre d'où provient l'erreur.

# Installez les bibliothèques 🤗 Transformers et 🤗 Datasets pour exécuter ce *notebook*.

# In[ ]:


get_ipython().system('pip install datasets transformers[sentencepiece]')


# In[ ]:


from datasets import load_dataset, load_metric
from transformers import (
    AutoTokenizer,
    TFAutoModelForSequenceClassification,
)

raw_datasets = load_dataset("glue", "mnli")

model_checkpoint = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)


def preprocess_function(examples):
    return tokenizer(examples["premise"], examples["hypothesis"], truncation=True)


tokenized_datasets = raw_datasets.map(preprocess_function, batched=True)

train_dataset = tokenized_datasets["train"].to_tf_dataset(
    columns=["input_ids", "labels"], batch_size=16, shuffle=True
)

validation_dataset = tokenized_datasets["validation_matched"].to_tf_dataset(
    columns=["input_ids", "labels"], batch_size=16, shuffle=True
)

model = TFAutoModelForSequenceClassification.from_pretrained(model_checkpoint)

model.compile(loss="sparse_categorical_crossentropy", optimizer="adam")

model.fit(train_dataset)


# In[ ]:


for batch in train_dataset:
    break


# In[ ]:


model.compile(optimizer="adam")


# In[ ]:


model(batch)


# In[ ]:


model = TFAutoModelForSequenceClassification.from_pretrained(model_checkpoint)
model(batch)


# In[ ]:


import numpy as np

loss = model(batch).loss.numpy()
indices = np.flatnonzero(np.isnan(loss))
indices


# In[ ]:


input_ids = batch["input_ids"].numpy()
input_ids[indices]


# In[ ]:


model.config.num_labels


# In[ ]:


from tensorflow.keras.optimizers import Adam

model = TFAutoModelForSequenceClassification.from_pretrained(model_checkpoint)
model.compile(optimizer=Adam(5e-5))


# In[ ]:


model.fit(train_dataset)


# In[ ]:


input_ids = batch["input_ids"].numpy()
tokenizer.decode(input_ids[0])


# In[ ]:


labels = batch["labels"].numpy()
label = labels[0]


# In[ ]:


for batch in train_dataset:
    break

# Assurez-vous que vous avez exécuté model.compile() et défini votre optimiseur,
# et vos pertes/métriques si vous les utilisez

model.fit(batch, epochs=20)