Text Classifier Fast API

text_classifier_api.py

from fastapi import FastAPI
import pickle
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB

app = FastAPI()

# Load the Tfidf and Naive Bayes models
tfidf = pickle.load(open("tf_idf.pkt", "rb"))
nb_model = pickle.load(open("toxicity_model.pkt", "rb"))

@app.post("/predict")
async def predict(text: str):
    # Transform the input text to Tfidf vectors
    text_tfidf = tfidf.transform([text]).toarray()
    
    # Predict the class of the input text
    prediction = nb_model.predict(text_tfidf)
    
    # Map the predicted class to a string
    class_name = "Toxic" if prediction == 1 else "Non-Toxic"
    
    # Return the prediction in a JSON response
    return {"text": text, "class": class_name}