vrkforever's picture
made changes to model
afcca73
raw
history blame
1.31 kB
import nltk
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import joblib
# Download necessary NLTK resources
nltk.download('wordnet', quiet=True)
nltk.download('stopwords', quiet=True)
# Initialize lemmatizer
lemmatizer = WordNetLemmatizer()
# Load the trained model
model = joblib.load('disaster_classification_model.joblib')
def improved_preprocess(text):
text = text.lower()
text = ''.join([char for char in text if char not in string.punctuation])
words = text.split()
words = [lemmatizer.lemmatize(word) for word in words if word not in stopwords.words('english')]
return ' '.join(words)
app = FastAPI()
class TextRequest(BaseModel):
text: str
@app.post("/predict")
async def predict(request: TextRequest):
text = request.text
new_text_processed = [improved_preprocess(text)]
prediction = model.predict(new_text_processed)
result = "disaster" if prediction == 1 else "not"
return JSONResponse(content={"output": result})
@app.get("/")
async def root():
return {"message": "Welcome to the Disaster Classification API"}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)