linksDetection / app.py
ZealPyae's picture
Upload app.py
59613f1 verified
raw
history blame contribute delete
950 Bytes
import joblib
import re
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from fastapi import FastAPI
from pydantic import BaseModel
# Load the model and vectorizer
vectorizer = joblib.load("vectorizer.joblib")
model = joblib.load("naive_bayes_model.joblib")
app = FastAPI()
class URLInput(BaseModel):
url: str
def preprocess_url(url):
url = re.sub(r"http\S+", "", url)
url = re.sub(r"\d+", "", url)
url = re.sub(r"\W", " ", url)
url = url.lower()
return url
@app.post("/predict")
def predict_url(url_input: URLInput):
processed_url = preprocess_url(url_input.url)
vectorized_url = vectorizer.transform([processed_url])
prediction = model.predict(vectorized_url)
return {"prediction": prediction[0]}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)