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 'not_spam'

def predict_PIPE(text):
  result = pipe(text)[0]
  return f'''{'spam' if result['label']=='LABEL_1' else 'not_spam'}
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()