Spaces:
Sleeping
Sleeping
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) | |
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]} | |
async def root(): | |
return {"message": "Spam Classification API"} |