ntam0001's picture
Create app.py
677f706 verified
raw
history blame
1.04 kB
from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import nltk
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer
import re
nltk.download('stopwords')
app = FastAPI()
# Load the model pipeline
pipeline = joblib.load('spam_classifier_pipeline.joblib')
class EmailRequest(BaseModel):
subject: str
body: str
def preprocess_text(text):
text = text.lower()
text = re.sub(r'[^a-zA-Z\s]', '', text)
words = text.split()
stop_words = set(stopwords.words('english'))
words = [word for word in words if word not in stop_words]
stemmer = PorterStemmer()
words = [stemmer.stem(word) for word in words]
return ' '.join(words)
@app.post("/predict")
async def predict(email: EmailRequest):
processed_text = preprocess_text(email.subject + ' ' + email.body)
prediction = pipeline.predict([processed_text])[0]
return {'prediction': ['ham', 'not_spam', 'spam'][prediction]}
@app.get("/")
async def root():
return {"message": "Spam Classification API"}