Update app.py

app.py

@@ -1,14 +1,17 @@
 #run the app
 #python -m streamlit run d:/NSFW/Project/test1.py
 import torch
-from transformers import BertTokenizer, BertForSequenceClassification
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+#from transformers import BertTokenizer, BertForSequenceClassification
 import math, keras_ocr
 # Initialize pipeline
 pipeline = None
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-model_2 = BertForSequenceClassification.from_pretrained("CustomModel")
-
-model_2.to('cpu')
+model_path="NSFW_text_classifier"
+#tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+#model_2 = BertForSequenceClassification.from_pretrained("CustomModel")
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForSequenceClassification.from_pretrained(model_path)
+#model_2.to('cpu')
 import streamlit as st
 
 def get_distance(predictions):
@@ -102,14 +105,17 @@ if uploaded_file is not None:
     sentance = ' '.join(ordered_preds)
     #st.write(sentance)
 
-    text =sentance
+    input_text =sentance
     print(text)
-    inputs = tokenizer(text,padding = True, truncation = True, return_tensors='pt').to('cpu')
+    """inputs = tokenizer(text,padding = True, truncation = True, return_tensors='pt').to('cpu')
     outputs = model_2(**inputs)
     predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    predictions = predictions.cpu().detach().numpy()
+    predictions = predictions.cpu().detach().numpy()"""
+    inputs = tokenizer(input_text, return_tensors="pt")
+    outputs = model(**inputs)
+    predictions = outputs.logits.softmax(dim=-1)
     print(predictions[0][0],predictions[0][1])
-    if predictions[0][0]>predictions[0][1]:
+    if predictions[0][1]>predictions[0][0]:
         print('safe')
         st.write('Safe for Work')
     else: