File size: 950 Bytes
59613f1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
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)