Recognises products by trade name or substance (Slovak inflection aware), queries LightRAG and shows the answer with the retrieved subgraph and ADC source texts.
139 lines
6.1 KiB
Python
139 lines
6.1 KiB
Python
"""Shared, offline-testable product recognition for the demo and evaluation."""
|
|
from __future__ import annotations
|
|
import re
|
|
import unicodedata
|
|
from collections import Counter, defaultdict
|
|
|
|
FORM_WORDS = {
|
|
"tbl", "cps", "sol", "crm", "gel", "ung", "inj", "inf", "plv", "sus", "sir", "gtt", "spr",
|
|
"flm", "obd", "eff", "ret", "pro", "por", "ijf", "dis", "lyo", "aer", "emp", "pst", "liq",
|
|
}
|
|
|
|
|
|
def normalize(text: str) -> str:
|
|
"""Lowercase, strip diacritics and punctuation (Slovak inflection is handled by prefix matching)."""
|
|
text = unicodedata.normalize("NFKD", text)
|
|
text = "".join(ch for ch in text if not unicodedata.combining(ch))
|
|
return re.sub(r"[^a-z0-9]+", " ", text.lower()).strip()
|
|
|
|
|
|
def build_aliases(_products: list[dict]) -> tuple[dict, dict, dict]:
|
|
"""Map normalized names to products.
|
|
|
|
Returns (aliases, by_prefix, by_url): aliases maps an alias to
|
|
{"kind": "substance"|"product", "label": str, "products": [index, ...]}.
|
|
"""
|
|
aliases: dict[str, dict] = {}
|
|
substance_products: dict[str, list[int]] = defaultdict(list)
|
|
brand_products: dict[str, list[int]] = defaultdict(list)
|
|
|
|
for i, product in enumerate(_products):
|
|
for substance in product["subs"]:
|
|
key = normalize(substance)
|
|
if len(key) >= 4:
|
|
substance_products[key].append(i)
|
|
tokens = normalize(product["name"]).split()
|
|
brand = []
|
|
for token in tokens:
|
|
if any(ch.isdigit() for ch in token) or token in FORM_WORDS:
|
|
if brand:
|
|
break
|
|
continue # leading "3M", "0 9 %" ... are not the brand
|
|
brand.append(token)
|
|
if brand and len(brand[0]) >= 3:
|
|
brand_products[brand[0]].append(i)
|
|
if len(brand) >= 2:
|
|
brand_products[" ".join(brand[:2])].append(i)
|
|
|
|
for key, ids in substance_products.items():
|
|
label = Counter(s for i in ids for s in _products[i]["subs"] if normalize(s) == key).most_common(1)[0][0]
|
|
aliases[key] = {"kind": "substance", "label": label, "products": ids}
|
|
for key, ids in brand_products.items():
|
|
if key in aliases:
|
|
continue
|
|
# A word shared by many unrelated products (e.g. "vitamin") is not a brand.
|
|
if len({tuple(_products[i]["subs"]) for i in ids}) > 6:
|
|
continue
|
|
aliases[key] = {"kind": "product", "label": key.title(), "products": ids}
|
|
|
|
by_prefix: dict[str, list[str]] = defaultdict(list)
|
|
for key in aliases:
|
|
by_prefix[key[:4]].append(key)
|
|
by_url = {p["url"]: p for p in _products if p["url"]}
|
|
return aliases, by_prefix, by_url
|
|
|
|
|
|
def detect_drugs(question: str, aliases: dict, by_prefix: dict) -> list[dict]:
|
|
"""Find products and active substances mentioned in the question.
|
|
|
|
Slovak nouns are inflected (warfarínom, ibuprofénu), so a question word
|
|
matches an alias it starts with, allowing a short ending.
|
|
"""
|
|
tokens = normalize(question).split()
|
|
found, used = {}, set()
|
|
for n in (2, 1):
|
|
for i in range(len(tokens) - n + 1):
|
|
if any(j in used for j in range(i, i + n)):
|
|
continue
|
|
phrase = " ".join(tokens[i:i + n])
|
|
for alias in by_prefix.get(phrase[:4], []):
|
|
if len(alias.split()) != n:
|
|
continue
|
|
exact = phrase == alias
|
|
inflected = len(alias) >= 5 and phrase.startswith(alias) and len(phrase) - len(alias) <= 4
|
|
if exact or inflected:
|
|
found[alias] = aliases[alias]
|
|
used.update(range(i, i + n))
|
|
break
|
|
return list(found.values())
|
|
|
|
|
|
def describe_match(match: dict, products: list[dict]) -> dict:
|
|
items = [products[i] for i in match["products"]]
|
|
substances = Counter(tuple(p["subs"]) for p in items if p["subs"]).most_common()
|
|
covered = [p for p in items if p["rep"]]
|
|
example = (covered or items)[0]
|
|
return {
|
|
"label": match["label"],
|
|
"kind": match["kind"],
|
|
"substances": list(substances[0][0]) if len(substances) == 1 else [],
|
|
"ambiguous": len(substances) > 1,
|
|
"substance_variants": [list(s) for s, _ in substances],
|
|
"count": len(items),
|
|
"in_graph": bool(covered),
|
|
"example": example,
|
|
}
|
|
|
|
|
|
BASE_PROMPT = (
|
|
"Odpovedaj po slovensky a iba v rozsahu dostupných dôkazov. "
|
|
"Rozlišuj indikácie, kontraindikácie, upozornenia a interakcie. "
|
|
"Neprítomnosť informácie neznamená neprítomnosť interakcie ani bezpečnosť kombinácie. "
|
|
"Uvádzaj zdroje k tvrdeniam, ak ich kontext poskytuje. Ak dôkaz chýba, priznaj to. "
|
|
"Nehádaj silu, liekovú formu ani zloženie nejednoznačnej značky. "
|
|
"Údaje z reprezentanta rovnakej účinnej látky nemusia pokrývať konkrétny prípravok."
|
|
)
|
|
|
|
# Reserve space for original ADC chunks after graph context and prompt overhead.
|
|
QUERY_SETTINGS = {"top_k":40,"chunk_top_k":10,"max_total_tokens":8000,
|
|
"max_entity_tokens":1000,"max_relation_tokens":1500,"enable_rerank":False}
|
|
|
|
|
|
def prepare_query(question, products, aliases=None, by_prefix=None):
|
|
if aliases is None or by_prefix is None:
|
|
aliases, by_prefix, _ = build_aliases(products)
|
|
matches = [describe_match(m, products) for m in detect_drugs(question, aliases, by_prefix)]
|
|
hints, notes = [], []
|
|
for match in matches:
|
|
if match["ambiguous"]:
|
|
notes.append(f"Názov {match['label']} je nejednoznačný; vyžiadaj presný názov a zloženie.")
|
|
elif match["substances"]:
|
|
hints.append(f"{match['label']} = účinná látka {', '.join(match['substances'])}")
|
|
if not match["in_graph"]:
|
|
ex = match["example"]
|
|
notes.append(f"Pre {ex['name']} nie je v indexe reprezentant s klinickými údajmi; "
|
|
"príčinou môže byť chýbajúci leták, parsovanie alebo výber dokumentov. "
|
|
f"Základný údaj ADC: {ex.get('ind','')}. Chýbajúce údaje neznamenajú bezpečnosť.")
|
|
query = question + ("\n\n(Lieky v otázke: " + "; ".join(hints) + ")" if hints else "")
|
|
return {"query":query,"matches":matches,"user_prompt":BASE_PROMPT + "\n" + "\n".join(notes)}
|