# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.6.9 64-bit (''bert'': virtualenv)'
#     metadata:
#       interpreter:
#         hash: eff88277fb73975c12d1396ee5027a873de6a4cc7778f3072fe3a9960a29a15a
#     name: 'Python 3.6.9 64-bit (''bert'': virtualenv)'
# ---

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import OneCycleLR
from transformers import BertTokenizerFast, BertModel
from transformers import AdamW
from utils.datasets import TextDataset
from utils.lr_finder import run_lr_finder
from models.bert import BertForClassification

import tqdm.notebook as tqdm

MODEL = "neuralmind/bert-base-portuguese-cased"

# # Loading tweet dataset

train_df = pd.read_csv("/home/kenzo/datasets/cleaned_tweetsentbr/train.tsv", sep="\t", names=["id", "label", "alfa", "text"], index_col=0)
test_df = pd.read_csv("/home/kenzo/datasets/cleaned_tweetsentbr/test.tsv", sep="\t", names=["id", "label", "alfa", "text"], index_col=0)

# +
# TODO: Exploração dos dados

# + tags=[]
tokenizer = BertTokenizerFast.from_pretrained(MODEL)
# -

train_ds = TextDataset.from_df(train_df, tokenizer, max_seq_len=128)
test_ds = TextDataset.from_df(test_df, tokenizer, max_seq_len=128)

# # Preparing model

# +
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from functools import partial

metrics = {
    "accuracy": accuracy_score,
    "precision": partial(precision_score, average="macro"),
    "recall": partial(recall_score, average="macro"),
    "f1": partial(f1_score, average="macro"),
}
# -

gpu = torch.device("cuda:2")

# + tags=[]
bert_model = BertModel.from_pretrained(MODEL)
# -

model = BertForClassification(bert_model, 3, metrics, device=gpu).cuda(gpu)

train_dl = DataLoader(train_ds, batch_size=64, shuffle=True)
test_dl = DataLoader(test_ds, batch_size=64, shuffle=True)

# # Finding a good learning rate

batches = len(train_dl)
epochs = 3
optimizer = AdamW(model.parameters(), lr=1e-7) 
# scheduler = OneCycleLR(optimizer, max_lr=5e-5, steps_per_epoch=batches, epochs=epochs)
criterion = nn.CrossEntropyLoss()

run_lr_finder(train_dl, model, optimizer, criterion, device=gpu)

scheduler = OneCycleLR(optimizer, max_lr=1e-5, steps_per_epoch=batches, epochs=epochs)

# + tags=[]
train, test = model.fit(epochs, train_dl, test_dl, criterion, optimizer, scheduler=scheduler)
# -

torch.save(model.state_dict(), "data/checkpoints/bert_tweetsent_br.ckpt")



StarcoderdataJupyterScriptsDedupFiltered