shreyankisiri's picture
Upload app.py
8350ed6 verified
raw
history blame contribute delete
952 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
import joblib
model = joblib.load("model.joblib")
tfidf_vectorizer = joblib.load("tfidf_vectorizer.joblib")
class TextInput(BaseModel):
text: str
app = FastAPI()
@app.post("/predict")
def predict(input: TextInput):
processed_text = preprocess_text(input.text)
text_tfidf = tfidf_vectorizer.transform([processed_text])
prediction = model.predict(text_tfidf)
return {"prediction": "Spam" if int(prediction[0]) == 0 else "Ham"}
def preprocess_text(text):
import re
from nltk.stem import WordNetLemmatizer
from nltk.corpus import stopwords
lemmatizer = WordNetLemmatizer()
stop_words = set(stopwords.words('english'))
text = re.sub('[^a-zA-Z]', ' ', text)
text = text.lower()
words = text.split()
words = [lemmatizer.lemmatize(word) for word in words if word not in stop_words]
return ' '.join(words)