app/models/graph/side_effect_head.py

"""Side-effect class predictions."""

from __future__ import annotations

from app.knowledge.side_effect_ontology import SIDE_EFFECT_TAGS


def predict_side_effects(drugs: list[str]) -> dict[str, float]:
    counts: dict[str, float] = {}
    for drug in drugs:
        for tag in SIDE_EFFECT_TAGS.get(drug, []):
            counts[tag] = counts.get(tag, 0.0) + 1.0
    total = sum(counts.values()) or 1.0
    return {k: v / total for k, v in counts.items()}