File size: 1,669 Bytes
eb4875c | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """
Gradio demo for the Watsonx Docs Type Classifier.
Loads the best trained model from models/ and serves predictions.
"""
from pathlib import Path
import gradio as gr
import joblib
import numpy as np
from sentence_transformers import SentenceTransformer
LABELS = ["conceptual", "how-to"]
model_name = (Path("models") / "best_model_name.txt").read_text().strip()
embedder = SentenceTransformer(model_name)
clf = joblib.load(Path("models") / "best_model.joblib")
def softmax(x):
e = np.exp(x - np.max(x))
return e / e.sum()
def predict(text: str) -> dict:
if not text.strip():
return {label: 0.0 for label in LABELS}
embedding = embedder.encode([text], convert_to_numpy=True)
if hasattr(clf, "predict_proba"):
probs = clf.predict_proba(embedding)[0]
else:
scores = clf.decision_function(embedding)[0]
# LinearSVC returns a scalar for binary; wrap in array
if np.ndim(scores) == 0:
scores = np.array([-scores, scores])
probs = softmax(scores)
return {label: float(p) for label, p in zip(LABELS, probs)}
demo = gr.Interface(
fn=predict,
inputs=gr.Textbox(
label="Document text (title + body)",
lines=8,
placeholder="Paste the title and opening text of a Watsonx documentation page here.",
),
outputs=gr.Label(num_top_classes=2, label="Predicted document type"),
title="Watsonx Docs Type Classifier",
description=(
"Predicts whether a Watsonx documentation page is **conceptual** or **how-to**. "
"Paste the page title and opening text below."
),
)
if __name__ == "__main__":
demo.launch(share=False)
|