fast-stager/STOP

1import os2import joblib3import gradio as gr4from fastapi import FastAPI, HTTPException5from pydantic import BaseModel, Field6from typing import List7 8CHECKPOINT_DIR = "checkpoints"9TFIDF_PATH = os.path.join(CHECKPOINT_DIR, "tfidf_vectorizer.pkl")10SVM_PATH = os.path.join(CHECKPOINT_DIR, "svm_stop_classifier.pkl")11 12LABEL_0 = "NOT_STOP"13LABEL_1 = "STOP"14 15tfidf_vectorizer = None16svm_model = None17 18try:19    print(f"Loading TFIDF Vectorizer from {TFIDF_PATH}...")20    tfidf_vectorizer = joblib.load(TFIDF_PATH)21    print(f"Loading SVM Model from {SVM_PATH}...")22    svm_model = joblib.load(SVM_PATH)23    print("Models loaded successfully.")24except FileNotFoundError as e:25    print(f"ERROR: Model file not found: {e}")26    raise RuntimeError(f"Failed to load required model files. Ensure 'checkpoint/' is correctly populated. Error: {e}")27 28app = FastAPI(29    title="STOP Classifier API",30    description="STOP/NOT_STOP text classification using Linear SVM. The main UI is at the root '/', while the API endpoints are at '/api-docs' and '/predict'.",31    version="1.0.0"32)33 34class PredictionRequest(BaseModel):35    texts: List[str] = Field(36        ...,37        description="A list of text strings to classify.",38        example=[39            "please discontinue all communication",40            "I will stop by the station after lunch"41        ]42    )43 44class PredictionResponse(BaseModel):45    text: str = Field(..., description="The input text.")46    prediction: str = Field(..., description="The predicted label (STOP or NOT_STOP).")47    probability_NOT_STOP: float = Field(..., description="Probability of NOT_STOP label.")48    probability_STOP: float = Field(..., description="Probability of STOP label.")49    inference_model: str = Field("SVM", description="The model used for inference.")50 51def predict_svm(texts: List[str]) -> List[PredictionResponse]:52    if not texts:53        return []54        55    vec = tfidf_vectorizer.transform(texts)56    probs = svm_model.predict_proba(vec)57    preds = svm_model.predict(vec)58 59    results = []60    for i, txt in enumerate(texts):61        pred_label = LABEL_1 if preds[i] == 1 else LABEL_062        results.append(PredictionResponse(63            text=txt,64            prediction=pred_label,65            probability_NOT_STOP=float(probs[i][0]),66            probability_STOP=float(probs[i][1]),67            inference_model="SVM"68        ))69    70    return results71 72@app.get("/health", status_code=200, tags=["API"])73def health_check():74    return {"status": "ok", "model_loaded": bool(svm_model)}75 76@app.post("/predict", response_model=List[PredictionResponse], tags=["API"])77async def post_predict(request: PredictionRequest):78    try:79        results = predict_svm(request.texts)80        return results81    except Exception as e:82        raise HTTPException(status_code=500, detail=f"Internal Server Error during POST prediction: {e}")83 84@app.get("/predict", response_model=PredictionResponse, tags=["API"])85async def get_predict(text: str):86    if not text.strip():87        raise HTTPException(status_code=400, detail="Text query parameter cannot be empty.")88        89    try:90        results = predict_svm([text])91        if not results:92            raise HTTPException(status_code=500, detail="Prediction returned empty result.")93        94        return results[0]95        96    except Exception as e:97        raise HTTPException(status_code=500, detail=f"Internal Server Error during GET prediction: {e}")98 99def gradio_interface_fn(text_input):100    """Interface function to be called by Gradio UI."""101    if not text_input or not text_input.strip():102        return "Please enter text for classification.", None103        104    try:105        result = predict_svm([text_input])[0]106        107        prediction_label = result.prediction108        109        prob_display = {110            LABEL_0: result.probability_NOT_STOP,  111            LABEL_1: result.probability_STOP       112        }113        114        return prediction_label, prob_display115        116    except Exception as e:117        return f"An error occurred: {str(e)}", None118    119ui = gr.Interface(120    fn=gradio_interface_fn,121    inputs=gr.Textbox(lines=2, placeholder="Enter a message to classify...", label="Input Text"),122    outputs=[123        gr.Label(label="Classification Result"), 124        gr.Label(label="Probabilities")125    ],126    title="STOP Classifier SVM",127    description="This is the user interface for the SVM model. The model classifies text as intended to end communication (STOP) or not (NOT_STOP). The API is available at the '/predict' endpoints."128)129 130app = gr.mount_gradio_app(app, ui, path="/")