"""Hugging Face Space demo for the embedding source classifier.""" from __future__ import annotations import os import inspect from functools import lru_cache from pathlib import Path import gradio as gr import joblib import numpy as np import yaml from huggingface_hub import hf_hub_download from sentence_transformers import SentenceTransformer DEFAULT_MODEL_REPO_ID = "birgermoell/saga-swedish-health-linear-probe" MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", DEFAULT_MODEL_REPO_ID) @lru_cache(maxsize=1) def load_assets(): model_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename="model.joblib") config_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename="resolved_config.yaml") classifier = joblib.load(model_path) config = yaml.safe_load(Path(config_path).read_text(encoding="utf-8")) embedder = load_embedder(config["model_id"]) return classifier, config, embedder def load_embedder(model_id: str): return RobustEmbedder(model_id) class RobustEmbedder: """Use SentenceTransformers first, then fall back to direct mean pooling.""" def __init__(self, model_id: str) -> None: self.model_id = model_id self.sentence_transformer = SentenceTransformer(model_id) self.fallback = None def encode(self, texts: list[str], normalize_embeddings: bool = True) -> np.ndarray: try: return self.sentence_transformer.encode(texts, normalize_embeddings=normalize_embeddings) except TypeError as exc: if "token_type_ids" not in str(exc): raise if self.fallback is None: self.fallback = MeanPoolingEmbedder(self.model_id) return self.fallback.encode(texts, normalize_embeddings=normalize_embeddings) class MeanPoolingEmbedder: """Fallback encoder for tokenizer/model pairs that reject token_type_ids.""" def __init__(self, model_id: str) -> None: import torch from transformers import AutoModel, AutoTokenizer self.torch = torch self.tokenizer = AutoTokenizer.from_pretrained(model_id) self.model = AutoModel.from_pretrained(model_id) self.model.eval() forward_signature = inspect.signature(self.model.forward) self.accepted_inputs = set(forward_signature.parameters) self.has_kwargs = any( parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in forward_signature.parameters.values() ) def encode(self, texts: list[str], normalize_embeddings: bool = True) -> np.ndarray: encoded = self.tokenizer( texts, padding=True, truncation=True, max_length=512, return_tensors="pt", ) if not self.has_kwargs: encoded = {key: value for key, value in encoded.items() if key in self.accepted_inputs} with self.torch.no_grad(): outputs = self.model(**encoded) token_embeddings = ( outputs.last_hidden_state if hasattr(outputs, "last_hidden_state") else outputs[0] ) attention_mask = encoded["attention_mask"].unsqueeze(-1).expand(token_embeddings.size()).float() pooled = (token_embeddings * attention_mask).sum(dim=1) / attention_mask.sum(dim=1).clamp(min=1e-9) if normalize_embeddings: pooled = self.torch.nn.functional.normalize(pooled, p=2, dim=1) return pooled.cpu().numpy().astype(np.float32) def classify(text: str): text = text.strip() if not text: return {}, "Enter a short health-information text." classifier, config, embedder = load_assets() embedding = embedder.encode([text], normalize_embeddings=True) prediction = classifier.predict(embedding)[0] probabilities = {} if hasattr(classifier, "predict_proba"): labels = _classifier_labels(classifier) scores = classifier.predict_proba(embedding)[0] probabilities = {label: float(score) for label, score in zip(labels, scores)} explanation = ( f"Predicted source: {prediction}\n\n" f"Embedding model: {config['model_id']}\n" "This demo is for an Information Retrieval assignment. It is not medical advice." ) return probabilities or {str(prediction): 1.0}, explanation def _classifier_labels(classifier): if hasattr(classifier, "classes_"): return [str(item) for item in classifier.classes_] if hasattr(classifier, "named_steps"): for step in reversed(classifier.named_steps.values()): if hasattr(step, "classes_"): return [str(item) for item in step.classes_] return [] examples = [ ( "You can often take care of a mild fever at home by resting, drinking fluids, " "and checking how you feel. Contact healthcare if the fever lasts several days, " "if you become very weak, or if a child has breathing difficulties." ), ( "Back pain is common and usually improves gradually. Keep moving if you can, " "avoid long periods of bed rest, and seek care urgently if you lose feeling in " "your legs or have trouble controlling your bladder or bowel." ), ( "Pollen allergy can cause sneezing, itchy eyes, and a blocked or runny nose. " "Non-prescription medicines may help, but contact healthcare if symptoms affect " "sleep, school, or work despite treatment." ), ( "The authority presents national indicators for stroke care, including treatment " "times, follow-up after discharge, and differences between regions. The results " "are intended to support quality improvement and equal care." ), ( "New national guidance recommends that municipalities and regions strengthen " "coordinated care for older people with multiple chronic conditions. The guidance " "summarizes evidence, priorities, and responsibilities for health services." ), ( "The annual statistical report describes trends in prescribed antibiotics, regional " "variation, and age-group differences. The data are used for monitoring public " "health and supporting national policy decisions." ), ( "The Medical Products Agency has received reports of a quality defect in a batch " "of tablets. Pharmacies and wholesalers should stop distribution while the company " "investigates whether patients need replacement packs." ), ( "A new safety communication reminds healthcare professionals to check liver function " "before and during treatment with the medicine. The product information will be " "updated with strengthened warnings and monitoring advice." ), ( "The agency approves a variation to the marketing authorisation after reviewing " "stability data, manufacturing documentation, and the proposed package leaflet. " "The decision concerns product quality and regulatory compliance." ), ] with gr.Blocks(title="Swedish Health Source Triage") as demo: gr.Markdown( """ # Swedish Health Source Triage Classify a short health-information text into a likely source family using sentence embeddings and a trained classifier. """ ) with gr.Row(): text = gr.Textbox(label="Text", lines=6, placeholder="Paste a short document or query...") with gr.Row(): label = gr.Label(label="Class probabilities") notes = gr.Textbox(label="Prediction notes", lines=5) gr.Examples(examples=examples, inputs=text) text.submit(classify, inputs=text, outputs=[label, notes]) gr.Button("Classify").click(classify, inputs=text, outputs=[label, notes]) if __name__ == "__main__": demo.launch()