Add files via upload

app.py

@@ -0,0 +1,188 @@
+import pandas as pd
+import numpy as np
+from torch.utils.data import Dataset
+import torch
+from transformers import AutoTokenizer
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from transformers import AutoModel, AdamW, get_cosine_schedule_with_warmup
+import torch.nn as nn
+import math
+from torchmetrics.functional.classification import auroc
+import torch.nn.functional as F
+import streamlit as st
+from transformers import pipeline
+
+
+
+class toxicity_dataset(Dataset):
+    def __init__(self,data_path,tokenizer,attributes,max_token_len= 128,sample = 1000):
+        self.data_path=data_path
+        self.tokenizer=tokenizer
+        self.attributes=attributes
+        self.max_token_len=max_token_len
+        self.sample=sample
+        self._prepare_data()
+    def _prepare_data(self):
+        data=pd.read_csv(self.data_path)
+        if self.sample is not None:
+            self.data=data.sample(self.sample,random_state=7)
+        else:
+            self.data=data
+    def __len__(self):
+        return(len(self.data))
+    def __getitem__(self,index):
+        item = self.data.iloc[index]
+        comment = str(item.comment_text)
+        attributes = torch.FloatTensor(item[self.attributes])
+        tokens = self.tokenizer.encode_plus(comment,add_special_tokens=True,return_tensors="pt",truncation=True,max_length=self.max_token_len,padding="max_length",return_attention_mask=True)
+        return{'input_ids':tokens.input_ids.flatten(),"attention_mask":tokens.attention_mask.flatten(),"labels":attributes}
+
+class Toxcity_Data_Module(pl.LightningDataModule):
+    def __init__(self,train_path,test_path,attributes,batch_size = 16, max_token_len = 128, model_name="roberta-base"):
+        super().__init__()
+        self.train_path=train_path
+        self.test_path=test_path
+        self.attributes=attributes
+        self.batch_size=batch_size
+        self.max_token_len=max_token_len
+        self.model_name=model_name
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+    def setup(self, stage = None):
+        if stage in (None, "fit"):
+            self.train_dataset=toxicity_dataset(self.train_path,self.tokenizer,self.attributes)
+            self.test_dataset=toxicity_dataset(self.test_path,self.tokenizer,self.attributes, sample=None)
+        if stage == "predict":
+            self.val_dataset=toxicity_dataset(self.test_path,self.tokenizer,self.attributes)
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset,batch_size=self.batch_size,shuffle=True)
+    def val_dataloader(self):
+        return DataLoader(self.train_dataset,batch_size=self.batch_size,shuffle=False)
+    def predict_dataloader(self):
+        return DataLoader(self.test_dataset,batch_size=self.batch_size,shuffle=False)
+    
+class Toxic_Comment_Classifier(pl.LightningModule):
+  def __init__(self, config: dict):
+    super().__init__()
+    self.config = config
+    self.pretrained_model = AutoModel.from_pretrained(config['model_name'], return_dict = True)
+    self.hidden = torch.nn.Linear(self.pretrained_model.config.hidden_size, self.pretrained_model.config.hidden_size)
+    self.classifier = torch.nn.Linear(self.pretrained_model.config.hidden_size, self.config['n_labels'])
+    torch.nn.init.xavier_uniform_(self.classifier.weight)
+    self.loss_func = nn.BCEWithLogitsLoss(reduction='mean')
+    self.dropout = nn.Dropout()
+    
+  def forward(self, input_ids, attention_mask=None, labels=None):
+    # roberta layer
+    output = self.pretrained_model(input_ids=input_ids, attention_mask=attention_mask)
+    pooled_output = torch.mean(output.last_hidden_state, 1)
+    # final logits
+    pooled_output = self.dropout(pooled_output)
+    pooled_output = self.hidden(pooled_output)
+    pooled_output = F.relu(pooled_output)
+    pooled_output = self.dropout(pooled_output)
+    logits = self.classifier(pooled_output)
+    # calculate loss
+    loss = 0
+    if labels is not None:
+      loss = self.loss_func(logits.view(-1, self.config['n_labels']), labels.view(-1, self.config['n_labels']))
+    return loss, logits
+
+  def training_step(self, batch, batch_index):
+    loss, outputs = self(**batch)
+    self.log("train loss ", loss, prog_bar = True, logger=True)
+    return {"loss":loss, "predictions":outputs, "labels": batch["labels"]}
+
+  def validation_step(self, batch, batch_index):
+    loss, outputs = self(**batch)
+    self.log("validation loss ", loss, prog_bar = True, logger=True)
+    return {"val_loss": loss, "predictions":outputs, "labels": batch["labels"]}
+
+  def predict_step(self, batch, batch_index):
+    loss, outputs = self(**batch)
+    return outputs
+
+  def configure_optimizers(self):
+    optimizer = AdamW(self.parameters(), lr=self.config['lr'], weight_decay=self.config['w_decay'])
+    total_steps = self.config['train_size']/self.config['bs']
+    warmup_steps = math.floor(total_steps * self.config['warmup'])
+    warmup_steps = math.floor(total_steps * self.config['warmup'])
+    scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)
+    return [optimizer],[scheduler]
+
+
+
+def predict_raw_comments(model, dm, trainer):
+  #print("debug1")
+  predictions = trainer.predict(model,dm)
+  #print("debug2")
+  flattened_predictions = np.stack([torch.sigmoid(torch.Tensor(p)) for batch in predictions for p in batch])
+  #print("debug3")
+  return flattened_predictions
+
+
+
+def main():
+    # -- Creates Variables for Use of Model -- 
+    attributes=["toxic","severe_toxic","obscene","threat","insult","identity_hate"]
+    tokenizer=AutoTokenizer.from_pretrained("roberta-base")
+    toxic_comments_dataset=toxicity_dataset("data/train.csv",tokenizer,attributes)
+
+    toxicity_data_module=Toxcity_Data_Module("data/train.csv","data/test.csv",attributes)
+    toxicity_data_module.setup()
+    dataloader=toxicity_data_module.train_dataloader()
+
+    config = {
+        'model_name':"distilroberta-base",
+        'n_labels':len(attributes),
+        'bs':128,
+        'lr':1.5e-6,
+        'warmup':0.2,
+        "train_size":len(toxicity_data_module.train_dataloader()),
+        'w_decay':0.001,
+        'n_epochs':1
+    }
+
+    toxicity_data_module=Toxcity_Data_Module("data/train.csv","data/reduced_test.csv",attributes,batch_size=config['bs'])
+    toxicity_data_module.setup()
+
+    
+    trainer = pl.Trainer(max_epochs=config['n_epochs'],num_sanity_val_steps=50)
+
+    ## -- Creates Streamlit App --
+    st.title("Tweet Toxicity Classifier ")
+    st.header("Fine tuned model from roberta-base using PyTorch")
+    st.header("Jozef Janosko - CS 482, Milestone 3")
+
+    model_name = st.selectbox("Select Model...", ["Toxicity Classification Model"])
+
+    if st.button("Click to Load Data"):
+        if model_name=="Toxicity Classification Model":
+            model = torch.load("ToxicityClassificationModel.pt")
+            with st.spinner('Analyzing Text...'):
+                logits = predict_raw_comments(model,toxicity_data_module,trainer=trainer)
+            torch_logits = torch.from_numpy(logits)
+            probabilities = F.softmax(torch_logits, dim = -1).numpy()
+            inputs=pd.read_csv("data/reduced_test.csv")
+            data=[]
+            #print(inputs["comment_text"][0]," ",probabilities)
+            for i in range(len(probabilities)):
+                max_prob = 0
+                max_cat = 6
+                
+                prob=0
+                for j in range(6):
+                    prob=probabilities[i][j]
+                    if(prob >= max_prob): 
+                        max_prob = prob 
+                        max_cat = j
+                #print(inputs["comment_text"][i]," ",attributes[max_cat]," ",max_prob," ",probabilities[i])
+                data.append([inputs["comment_text"][i][0:16]+"...",attributes[max_cat],max_prob])
+            results_df=pd.DataFrame(data,columns=["Comment Text","Most Likely Classification","Classification Probability"])
+            st.table(data=results_df)
+
+    
+    
+
+if __name__ == '__main__' :
+    main()