spam-classifier / app.py
vishal2023's picture
Update app.py
880aa1f
raw
history blame
1.2 kB
import joblib
import string
import re
import nltk
nltk.download('stopwords')
from nltk.corpus import stopwords
from nltk.stem.porter import PorterStemmer
import gradio as gr
model = joblib.load('naive_bayes_model.sav')
tfidf = joblib.load('TfIdfVectorizer.sav')
pipe = joblib.load('pipeline.sav')
def predict_NB(text):
text = text.lower()
text = re.sub('[^A-Za-z]',' ',text)
text = text.translate(str.maketrans('','',string.punctuation))
text = ' '.join(word for word in text.split() if word not in stopwords.words('english'))
ps = PorterStemmer()
text = ' '.join([ps.stem(word) for word in text.split()])
X = tfidf.transform([text]).toarray()
return 'spam' if model.predict(X)[0] == 1 else 0
def predict_PIPE(text):
result = pipe(text)[0]
return f'''{'spam' if result['label']=='LABEL_1' else 0}
confidence : {result['score']}'''
def fn(model_choice, input):
if model_choice=="naive-bayes":
return predict_NB(input)
elif model_choice=="tiny-bert":
return predict_PIPE(input)
gr.Interface(fn, inputs = [gr.inputs.Dropdown(["naive-bayes", "tiny-bert"],default = 'naive-bayes'),'text'],
outputs = "text",
title = 'Spam Classifier').launch()