Spam Detection

import re, string
def clean(text: str) -> str:
    text = text.lower()
    text = re.sub(r"http\S+", " URL ", text)
    text = re.sub(r"\d+", " NUM ", text)
    text = text.translate(str.maketrans("", "", string.punctuation))
    return " ".join(text.split())

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.metrics import classification_report

pipe = Pipeline([
    ("tfidf", TfidfVectorizer(ngram_range=(1,2), min_df=2, max_features=20000)),
    ("clf",   LogisticRegression(max_iter=1000, class_weight="balanced")),
])
pipe.fit(X_train, y_train)
print(classification_report(y_test, pipe.predict(X_test)))

from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
import datasets

ds = datasets.load_dataset("sms_spam")
tok = AutoTokenizer.from_pretrained("distilbert-base-uncased")
def enc(b): return tok(b["sms"], truncation=True, padding="max_length", max_length=64)
ds = ds.map(enc, batched=True).rename_column("label","labels")

model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)
args = TrainingArguments("out", num_train_epochs=3, per_device_train_batch_size=32,
                         evaluation_strategy="epoch", learning_rate=2e-5)
Trainer(model=model, args=args, train_dataset=ds["train"], eval_dataset=ds["train"]).train()

from fastapi import FastAPI
from pydantic import BaseModel
import joblib
pipe = joblib.load("spam_pipe.pkl")
app = FastAPI()

class Msg(BaseModel):
    text: str

@app.post("/predict")
def predict(m: Msg):
    p = pipe.predict_proba([m.text])[0,1]
    return {"spam_probability": float(p), "is_spam": bool(p > 0.5)}