snaramirez872/Finetuning-Toxicity-Model

1import torch2import torch.nn as TNN3import pandas as pd4from tqdm import tqdm5from torch.utils.data import Dataset as set, DataLoader as DL6from torch import cuda7import streamlit as st8from transformers import BertTokenizer as BT, BertModel as BM9 10device = 'cuda' if cuda.is_available() else 'cpu'11 12# Defined variables for later use13MAX_LEN = 12814TRAIN_BATCH_SIZE = 415VALID_BATCH_SIZE = 416LEARNING_RATE = 5e-0517 18modName = 'bert-base-uncased' # Pre-trained model19categories = ['toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate'] # Labels20 21data = pd.read_csv('./train.csv')22data.drop(['id'], inplace=True, axis=1)23 24new = pd.DataFrame()25new['text'] = data['comment_text']26new['labels'] = data.iloc[:,1].values.tolist()27 28tokenizer = BT.from_pretrained(modName, truncation=True, do_lower_case=True)29 30class MultiLabelDataset(set):31    def __init__(self, df, tokenizer, max_len):32        self.tokenizer = tokenizer33        self.data = df34        self.text = df.text35        self.targets = self.data.labels36        self.max_len = max_len37 38    def __len__(self):39        return len(self.targets)40 41    def __getitem__(self, idx):42        text = str(self.text[idx])43        text = " ".join(text.split())44 45        ins = self.tokenizer.encode_plus(46            text,47            None,48            add_special_tokens=True,49            max_length=self.max_len,50            pad_to_max_length=True,51            return_token_type_ids=True52        )53        input_ids = ins['input_ids']54        attention_mask = ins['attention_mask']55        token_type_ids = ins["token_type_ids"]56 57        #st.write("Input Keys: ", ins.keys()) # was used for debugging58        return {59            'input_ids': torch.tensor(input_ids, dtype=torch.long),60            'attention_mask': torch.tensor(attention_mask, dtype=torch.long),61            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),62            'targets': torch.tensor(self.targets[idx], dtype=torch.float)63        }64 65trainSize = 0.866trainData = new.sample(frac=trainSize,random_state=200)67testData = new.drop(trainData.index).reset_index(drop=True)68trainData = trainData.reset_index(drop=True)69 70trainSet = MultiLabelDataset(trainData, tokenizer, MAX_LEN)71testSet = MultiLabelDataset(testData, tokenizer, MAX_LEN)72 73training_loader = DL(trainSet, batch_size=TRAIN_BATCH_SIZE, shuffle=True)74testing_loader = DL(testSet, batch_size=VALID_BATCH_SIZE, shuffle=True)75 76# neural network77class BERTClass(TNN.Module):78    def __init__(self):79        super(BERTClass, self).__init__()80        self.l1 = BM.from_pretrained(modName)81        self.pre_classifier = TNN.Linear(768, 768)82        self.dropout = TNN.Dropout(0.1)83        self.classifier = TNN.Linear(768, 6)84 85    def forward(self, input_ids, attention_mask, token_type_ids):86        out = self.l1(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)87        hidden_state = out[0]88        po = hidden_state[:, 0]89        po = self.pre_classifier(po)90        po = TNN.Tanh()(po)91        po = self.dropout(po)92        outs = self.classifier(po)93        return outs94 95mod = BERTClass()96mod.to(device)97 98# Loss function and Optimizer99def lossFN(outs, targets):100    targets = targets.unsqueeze(1).expand_as(outs)101    return TNN.BCEWithLogitsLoss()(outs, targets)102 103opt = torch.optim.Adam(mod.parameters(), lr=LEARNING_RATE)104 105# Training and Finetuning106def train(mod, training_loader):107    mod.train()108    for _, data in tqdm(enumerate(training_loader, 0)):109        input_ids = data['input_ids'].to(device, dtype=torch.long)110        attention_mask = data['attention_mask'].to(device, dtype=torch.long)111        token_type_ids = data['token_type_ids'].to(device, dtype=torch.long)112        targets = data['targets'].to(device, dtype=torch.float)113 114        outs = mod(input_ids, attention_mask, token_type_ids)115 116        opt.zero_grad()117        loss = lossFN(outs, targets)118        loss.backward()119        opt.step()120 121# StreamLit Table of Results122st.title("Finetuned Model for Toxicity")123st.subheader("Model: bert-base-uncased")124 125def predict(tweets):126    mod.eval()127    res = []128    with torch.no_grad():129        for ins in tweets:130            outs = mod(input_ids=ins['input_ids'].to(device), attention_mask=ins['attention_mask'].to(device), token_type_ids=ins['token_type_ids'].to(device))131            probs = torch.softmax(outs[0], dim=-1)132            preds = torch.argmax(probs, dim=-1)133            for i in range(len(tweets)):134                res.append({'TWEETS': tweets, 'LABEL': preds[i].item(), 'PROBABILITY': probs[i][preds[i].item()].item()})135    return res136 137res = predict(testing_loader)138st.table(res) # table139