Update app.py

app.py

@@ -1,54 +1,34 @@
-from transformers import pipeline
+import pandas as pd
 import torch
 import torch.nn.functional as TF
 import streamlit as st
 
-model_name = "RoBERTa"
+option = st.selectbox("Select a text analysis model:", ("BERT", "Fine-tuned BERT"))
 
-classifier = pipeline("sentiment-analysis")
-defaultTxt = "I hate you cancerous insects so much"
-result = classifier(defaultTxt)
-st.write(result)
-
-if (option == "RoBERTa"):
-    tokenizerPath = "s-nlp/roberta_toxicity_classifier"
-    modelPath = "s-nlp/roberta_toxicity_classifier"
-    neutralIndex = 0
-    toxicIndex = 1
-elif (option == "DistilBERT"):
-    tokenizerPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
-    modelPath = "citizenlab/distilbert-base-multilingual-cased-toxicity"
-    neutralIndex = 1
-    toxicIndex = 0
-elif (option == "XLM-RoBERTa"):
-    tokenizerPath = "unitary/multilingual-toxic-xlm-roberta"
-    modelPath = "unitary/multilingual-toxic-xlm-roberta"
-    neutralIndex = 1
-    toxicIndex = 0
+bert_path = "bert-base-uncased"
+if (option == "BERT"):
+    
+    tokenizer = AutoTokenizer.from_pretrained(bert_path)
+    model = BertForSequenceClassification.from_pretrained(bert_path, num_labels=6)
 else:
-    tokenizerPath = "s-nlp/roberta_toxicity_classifier"
-    modelPath = "s-nlp/roberta_toxicity_classifier"
-    neutralIndex = 0
-    toxicIndex = 1
-
-tokenizer = AutoTokenizer.from_pretrained(tokenizerPath)
-model = AutoModelForSequenceClassification.from_pretrained(modelPath)
+    
 
-tokens = tokenizer.tokenize(input_text)
-token_ids = tokenizer.convert_tokens_to_ids(tokens)
-input_ids = tokenizer(input_text)
+tweets_raw = pd.read_csv("train.csv", nrows=20)
 
-batch = tokenizer(X_train, padding=True, truncation=True, max_length=512, return_tensors="pt")
 
-with torch.no_grad():
-    outputs = model(**batch)
-    predictions = TF.softmax(outputs.logits, dim=1)
-    labels = torch.argmax(predictions, dim=1)
-    labels = [model.config.id2label[label_id] for label_id in labels.tolist()]
+# Run encoding through model to get classification output.
+encoding = tokenizer.encode(txt, return_tensors='pt')
+result = model(encoding)
 
-save_directory = "saved"
-tokenizer.save_pretrained(save_directory)
-model.save_pretrained(save_directory)
+# Transform logit to get probabilities.
+if (result.logits.size(dim=1) < 2):
+    pad = (0, 1)
+    result.logits = nn.functional.pad(result.logits, pad, "constant", 0)
+prediction = nn.functional.softmax(result.logits, dim=-1)
+neutralProb = prediction.data[0][neutralIndex]
+toxicProb = prediction.data[0][toxicIndex]
 
-tokenizer = AutoTokenizer.from_pretrained(save_directory)
-model = AutoModelForSequenceClassification.from_pretrained(save_directory)
+# Write results
+st.write("Classification Probabilities")
+st.write(f"{neutralProb:.4f} - NEUTRAL")
+st.write(f"{toxicProb:.4f} - TOXIC")