new app

app.py

@@ -1,4 +1,10 @@
+import re
 import gradio as gr
+from collections import defaultdict
+from ai4data import extract_from_text, deduplicate_mentions
+import html
+from typing import List
+
 import os
 import subprocess
 
@@ -12,318 +18,491 @@ if GH_TOKEN:
         check=True
     )
 else:
-    print("⚠️ GH_TOKEN not found. Private ai4data_use will NOT be installed.")
+    print("GH_TOKEN not found. Private ai4data_use will NOT be installed.")
 
-from ai4data import extract_from_text
 DATA_MODEL_ID = "rafmacalaba/datause-extraction-v3-finetuned"
-DATASET_COLORS = {
-    "named": "#ff6b6b",
-    "unnamed": "#51cf66",
-    "vague": "#fcc419",
-    "invalid": "#868e96"
-}
-
-RELATION_COLORS = {
-    "acronym": "#4dabf7",                 # REL_ACR
-    "author": "#f06595",                  # REL_AUT
-    "data description": "#e599f7",        # REL_DSC
-    "data geography": "#20c997",          # REL_GEO
-    "data type": "#339af0",               # REL_TYP
-    "publication year": "#f783ac",        # REL_PY
-    "publisher": "#4dabf7",               # REL_PUB
-    "reference population": "#8e44ad",    # REL_POP
-    "reference year": "#ffd43b",          # REL_RY
-    # Removed on purpose — we are filtering these out
-    # "usage context": "#e8590c"
-}
-
-RELATION_DISPLAY = {
-    "data geography": "geography",
-    "data description": "description",
-}
-
-
-def overlaps(a, b):
-    # spans overlap if they share any character position
-    return not (a["end"] <= b["start"] or a["start"] >= b["end"])
-
-
-# def collect_spans(res):
-    base = res[0]["text"]
-
-    dataset_spans = []
-    relation_spans = []
-    usage_context_map = {}
-    # --- Collect DATASETS (always keep) ---
-    for item in res:
+
+def tokenize_text(text: str) -> List[str]:
+    """Tokenize the input text into a list of tokens (words, punctuation)."""
+    return re.findall(r'\w+(?:[-_]\w+)*|\S', text)
+
+def truncate_text(text: str, max_tokens: int = 300) -> str:
+    """
+    Tokenize the text and truncate to the first `max_tokens` tokens,
+    rejoining them into a valid string.
+    """
+    tokens = tokenize_text(text)
+    if len(tokens) <= max_tokens:
+        return text.strip()
+    truncated = tokens[:max_tokens]
+    return " ".join(truncated).strip() + " ..."
+
+def process_text(text, distance_threshold=1000, short_threshold=100, debug=False, precomputed_results=None):
+    if not text.strip():
+        return [], []
+
+    # Use precomputed results if provided
+    results = precomputed_results if precomputed_results is not None else extract_from_text(text)
+    if not results:
+        return [], []
+
+    base_text = results[0]["text"].strip()
+    from collections import defaultdict
+    import re
+
+    entity_map = defaultdict(list)
+    dataset_spans = []  # (start, end, text, label)
+    all_relations = []
+    absorbed_vague_datasets = set()
+
+    # --- 1️⃣ Collect dataset spans ---
+    for item in results:
         ds = item.get("datasets")
         if ds:
-            dataset_spans.append({
-                "start": ds["start"],
-                "end": ds["end"],
-                "label": ds["label"].lower(),
-                "kind": "dataset",
-                "text": base[ds["start"]:ds["end"]]
-            })
-
-    # --- Collect RELATIONS (but skip usage context) ---
-    for item in res:
-        for r in item.get("relations", []):
-            if r["relation"].lower() == "usage context":   # remove usage context entirely
+            start, end = ds["start"], ds["end"]
+            ds_label = ds["label"]
+            ds_text = ds["text"]
+            dataset_spans.append((start, end, ds_text, ds_label))
+
+    # --- 2️⃣ Collect relations (keep for tree even if overlapping) ---
+    seen_descriptions = set()
+    for item in results:
+        for rel in item.get("relations", []):
+            start, end = rel["start"], rel["end"]
+            src = rel["source"]
+            relation = rel["relation"].strip().lower()
+            target_text = base_text[start:end].strip()
+
+            # Skip implausible year relations
+            if relation in {"reference year", "publication year"} and not re.search(r"\b\d{4}\b", target_text):
                 continue
-            if r.get("score", 0) < 0.6:
-                continue  # ignore low-confidence relations
-            relation_spans.append({
-                "start": r["start"],
-                "end": r["end"],
-                "label": r["relation"].lower(),
-                "kind": "relation",
-                "text": base[r["start"]:r["end"]],
-                "source": r['source']
-            })
-
-    # --- Priority Rule: Dataset spans win over relation overlaps ---
-    def overlaps(a, b):
-        return not (a["end"] <= b["start"] or a["start"] >= b["end"])
-
-    dataset_spans = sorted(dataset_spans, key=lambda s: s["start"])
-
-    filtered_relations = []
-    for rel in relation_spans:
-        if any(overlaps(rel, ds) for ds in dataset_spans):
+
+            limit = short_threshold if relation in {"data description", "data type"} else distance_threshold
+            linked_datasets = []
+
+            # --- Find nearby datasets within threshold ---
+            if dataset_spans:
+                for ds_start, ds_end, ds_text, _ in dataset_spans:
+                    dist = abs(ds_start - start)
+                    if dist <= limit:
+                        linked_datasets.append(ds_text)
+
+            # --- If none found, apply shared-relation logic ---
+            if not linked_datasets:
+                # Shared or global relations — always apply to all datasets
+                if relation in {"publisher", "reference year", "publication year", "usage context"} and dataset_spans:
+                    linked_datasets = [d[2] for d in dataset_spans]
+                else:
+                    continue  # skip too-distant or unrelated relation
+
+            # --- Supersession rule: vague dataset near named one ---
+            if relation in {"data type"}:
+                for ds_start, ds_end, ds_text, ds_label in dataset_spans:
+                    # only absorb if relation's source explicitly mentions this dataset
+                    if ds_label == "vague" and ds_text in rel["source"] and abs(ds_start - start) <= short_threshold:
+                        absorbed_vague_datasets.add(ds_text)
+
+            # --- Deduplication for data description ---
+            if relation == "data description":
+                key = target_text.lower()
+                if key in seen_descriptions:
+                    continue
+                seen_descriptions.add(key)
+
+            # --- Store relation for all linked datasets ---
+            for ds_text in linked_datasets:
+                relation_entry = {
+                    "start": start,
+                    "end": end,
+                    "relation": relation,
+                    "src": ds_text,
+                }
+
+                # only store text value for usage context
+                if relation == "usage context":
+                    relation_entry["value"] = target_text or ""
+
+                all_relations.append(relation_entry)
+                if debug:
+                    print(f"Linked {relation} → {ds_text}")
+
+    # --- 3️⃣ Add dataset highlights (priority layer) ---
+    for ds_start, ds_end, ds_text, ds_label in dataset_spans:
+        if ds_label == "vague" and ds_text in absorbed_vague_datasets:
             continue
-        filtered_relations.append(rel)
+        entity_map[(ds_start, ds_end)].append(f"dataset:{ds_label}")
+
+    # --- 4️⃣ Add relation spans, skip those overlapping datasets ---
+    relation_priority = [
+        "acronym", "publisher", "data type", "data description",
+        "reference year", "publication year", "data geography", "usage context"
+    ]
+
+    relation_map = defaultdict(list)
+    for rel in all_relations:
+        key = (rel["start"], rel["end"])
+        relation_map[key].append(rel["relation"])
+
+    for (start, end), rels in relation_map.items():
+        # --- improved overlap: skip if relation is fully or partially inside a dataset span ---
+        overlaps_dataset = False
+        for ds_start, ds_end, _, _ in dataset_spans:
+            if (start >= ds_start and end <= ds_end) or (start < ds_end and end > ds_start):
+                overlaps_dataset = True
+                break
+
+        if overlaps_dataset:
+            continue  # keep for tree, not highlight
+
+        # choose top relation if multiple overlap on same span
+        if len(rels) > 1:
+            rels.sort(key=lambda r: relation_priority.index(r) if r in relation_priority else len(relation_priority))
+        top_relation = rels[0]
+        entity_map[(start, end)].append(top_relation)
+
+    # --- 5️⃣ Rebuild final text output for Gradio highlight ---
+    sorted_spans = sorted(entity_map.items(), key=lambda x: x[0][0])
+    output = []
+    last_idx = 0
+    for (start, end), labels in sorted_spans:
+        if start > last_idx:
+            output.append((base_text[last_idx:start], None))
+        substring = base_text[start:end]
+        label = ", ".join(labels)
+        output.append((substring, label))
+        last_idx = end
+
+    if last_idx < len(base_text):
+        output.append((base_text[last_idx:], None))
+
+    return output, results
+
+
+
+# --- 🌍 Long, realistic research examples ---
+example_texts = [
+    # 1️⃣ Ghana mining + household surveys
+    """Introduction The mining sector in Africa is growing rapidly and remains the main recipient of foreign direct investment ( World Bank 2011 ). The welfare effects of this sector are not well understood, although a literature has recently developed around this question. The main contribution of this paper is to shed light on the welfare effects of gold mining in a detailed, in-depth country study of Ghana, a country with a long tradition of gold mining and a recent, large expansion in capital-intensive and industrial-scale production. We use two complementary geocoded household data sets to analyze outcomes in Ghana: the Demographic and Health Survey ( DHS ) and the Ghana Living Standard Survey ( GLSS ), both of which provide information on household welfare, education, and demographic outcomes. The empirical analysis relies on district-level spatial models and difference-in-differences estimation to capture local spillover effects. We also examine how outcomes vary by proximity to large mines and across survey years from 2005 to 2015, controlling for baseline characteristics and mining expansion phases.""",
+
+    # 2️⃣ WDI, IEA, and energy transition
+    """This study analyzes labor and household data in the context of Afghan refugees in Iran. The national household surveys in Iran have good coverage of them. 
+The Labor Force Survey (LFS) has been conducted in the middle of each quarter since spring 2005 and is the primary source of employment statistics in Iran. 
+The LFS covers a wide range of topics including household members’ demographic and employment status, such as education, migration, working hours, industry, occupation, and experience (but not wage and income). 
+Importantly for this research, the nationality of each household member is also inquired in this survey. The sampling of LFS is on a rotating panel basis in the sense that each household is sampled in two consecutive seasons of two consecutive years. 
+This feature enables us to observe each individual’s change in employment status and compare it between two communities. Table 1 lists the LFS rounds with the available sample of Afghan refugees. 
+HEIS (Household Expenditure and Income Survey) has been conducted annually since 1963, but its raw data is available from 1984 onwards.""",
+
+    # 3️⃣ DHS and MICS in maternal and child health
+   """Observations are also spatially identified at the municipality level, but here we focus on variation in the Venezuelan share of the population at the province level, of which there are 196, as these are best representative of local labor markets. The Latin American Public Opinion Project (LAPOP) is an opinion survey conducted bi-annually in all countries in Latin America and designed to be representative of urban populations. This was fielded in Peru in 2010, 2012, 2014, 2017 and 2019 and consists of about 2,000 observations from mostly urban areas."""
+,
+
+    # 4️⃣ UNHCR, ProGres, and displacement data
+    """Such null effects of refugee migration on native attitudes in host communities are also identified by Zhou et al. (2021) for the case of Sudanese refugee immigration to Uganda. Our analysis relies on data from the following sources: the Encuesta Dirigida a la Población Venezolana que Reside en el País (ENPOVE), which is a specialized survey of Venezuelans living in Peru conducted by the National Institute of Statistics (INEI) in December 2018. The survey covers five main urban areas in the country where Venezuelan immigrants were most likely to be present.""",
+]
+
+def make_tree_from_highlight(highlight_output, max_distance=250):
+    """
+    Build dataset–relation tree using simplified labels from process_text().
+    Proximity-based grouping: attach relations to nearest dataset mention.
+    """
+    if not highlight_output:
+        return "<i>No relations found.</i>"
+
+    # Collect dataset mentions and relation mentions
+    datasets = []
+    relations = []
+    for i, (segment, label) in enumerate(highlight_output):
+        if not label:
+            continue
+        labels = [lbl.strip() for lbl in label.split(",")]
+        for lbl in labels:
+            if lbl.startswith("dataset:"):
+                datasets.append({
+                    "idx": i,
+                    "name": segment.strip(),
+                    "label": lbl.split(":")[1].strip()
+                })
+            elif lbl in {
+                "acronym", "publisher", "data type", "data description",
+                "reference year", "publication year", "data geography", "usage context"
+            }:
+                # ✅ Store 'target' and 'value' separately — only for usage context
+                rel_entry = {"idx": i, "type": lbl, "target": segment.strip()}
+                if lbl == "usage context":
+                    rel_entry["value"] = segment.strip()  # target text = value
+                relations.append(rel_entry)
+
+    if not datasets and not relations:
+        return "<i>No relations found.</i>"
+
+    # Build a dict to hold dataset → {relation_type → [targets]}
+    tree = {ds["name"]: {"_dataset_label": ds["label"]} for ds in datasets}
+
+    # Attach relations to nearest dataset mention (proximity-based)
+    for rel in relations:
+        if not datasets:
+            continue
+        nearest = min(datasets, key=lambda ds: abs(ds["idx"] - rel["idx"]))
+        if abs(nearest["idx"] - rel["idx"]) <= max_distance:
+            tree.setdefault(nearest["name"], {"_dataset_label": nearest["label"]})
+            rel_type = rel["type"]
+            # ✅ store 'value' separately for usage context
+            if rel_type == "usage context":
+                tree[nearest["name"]].setdefault(rel_type, set()).add(rel.get("value", ""))
+            else:
+                tree[nearest["name"]].setdefault(rel_type, set()).add(rel["target"])
+
+    # --- Render HTML ---
+    html_out = ["<div style='font-family:monospace; font-size:0.9em; line-height:1.5;'>"]
+    for ds_name, rels in sorted(tree.items()):
+        ds_label = rels.get("_dataset_label")
+        ds_badge = f" <span style='color:gray;'>[{html.escape(ds_label)}]</span>" if ds_label else ""
+        html_out.append(
+            f"<details open><summary><b style='color:#4da6ff;'>{html.escape(ds_name)}</b>{ds_badge}</summary><ul>"
+        )
+        for rel_type, targets in sorted(rels.items()):
+            if rel_type == "_dataset_label":
+                continue
+            for t in sorted(targets):
+                # ✅ Conditional rendering for usage context
+                if rel_type == "usage context":
+                    html_out.append(
+                        f"<li><span style='color:#80c904;'>{html.escape(rel_type)}</span>: "
+                        f"<i>{html.escape(t)}</i></li>"
+                    )
+                else:
+                    html_out.append(
+                        f"<li><span style='color:#80c904;'>{html.escape(rel_type)}</span>: {html.escape(t)}</li>"
+                    )
+        html_out.append("</ul></details>")
+    html_out.append("</div>")
+
+    return "\n".join(html_out)
+
+import html
+
+def make_relation_tree_from_results(highlight_output, results):
+    """
+    Build relation tree directly from model results,
+    ensuring all relations (even those not highlighted) are captured.
+    """
+    if not results:
+        return "<i>No extracted relations found.</i>"
+
+    # Collect dataset mentions and their relations
+    datasets = []
+    tree = {}
 
-    # --- Deduplicate spans ---
-    uniq = {}
-    for span in dataset_spans + filtered_relations:
-        key = (span["start"], span["end"], span["label"], span["kind"])
-        uniq[key] = span  # last wins (order doesn't matter because datasets sorted first)
+    for item in results:
+        ds = item.get("datasets")
+        rels = item.get("relations", [])
+        if not ds:
+            continue
 
-    spans = list(uniq.values())
-    spans = sorted(spans, key=lambda s: s["start"])
+        ds_name = ds.get("text", "").strip()
+        ds_label = ds.get("label", "")
+        if not ds_name:
+            continue
 
-    return spans
+        # Initialize dataset node
+        if ds_name not in tree:
+            tree[ds_name] = {"_dataset_label": ds_label}
 
-def collect_spans(res):
-    base = res[0]["text"]
+        # Collect all relations
+        for rel in rels:
+            rel_type = rel.get("relation", "").strip().lower()
+            target = rel.get("target", "").strip()
 
-    dataset_spans = []
-    relation_spans = []
+            # Skip empty or malformed relations
+            if not rel_type or not target:
+                continue
 
-    # We store usage context *indexed by item*, not by string
-    usage_context_for_item = {}
+            # Special handling for usage context (value-only)
+            if rel_type == "usage context" and "value" in rel:
+                target = rel["value"]
 
-    # --- First: capture usage context by item index ---
-    for idx, item in enumerate(res):
-        for r in item.get("relations", []):
-            if r["relation"].lower() == "usage context":
-                usage_context_for_item[idx] = r.get("target", "").lower()
+            # Attach relation under this dataset
+            tree[ds_name].setdefault(rel_type, set()).add(target)
 
-    # --- Collect DATASETS and attach usage context correctly ---
-    for idx, item in enumerate(res):
-        ds = item.get("datasets")
-        if ds:
-            usage = usage_context_for_item.get(idx, None)
-            dataset_spans.append({
-                "start": ds["start"],
-                "end": ds["end"],
-                "label": ds["label"].lower(),
-                "kind": "dataset",
-                "text": base[ds["start"]:ds["end"]],
-                "usage": usage     # ✅ always aligned — no string matching needed
-            })
-
-    # --- Collect RELATIONS (excluding usage context & low confidence) ---
-    for idx, item in enumerate(res):
-        for r in item.get("relations", []):
-            rel_label = r["relation"].lower()
-            if rel_label == "usage context":
-                continue
-            if r.get("score", 0) < 0.6:
+    # Render HTML
+    html_out = ["<div style='font-family:monospace; font-size:0.9em; line-height:1.5;'>"]
+    for ds_name, rels in sorted(tree.items()):
+        ds_label = rels.get("_dataset_label", "")
+        ds_badge = f" <span style='color:gray;'>[{html.escape(ds_label)}]</span>" if ds_label else ""
+        html_out.append(
+            f"<details open><summary><b style='color:#4da6ff;'>{html.escape(ds_name)}</b>{ds_badge}</summary><ul>"
+        )
+        for rel_type, targets in sorted(rels.items()):
+            if rel_type == "_dataset_label":
                 continue
-            relation_spans.append({
-                "start": r["start"],
-                "end": r["end"],
-                "label": rel_label,
-                "kind": "relation",
-                "text": base[r["start"]:r["end"]],
-                "source": r['source'],
-                "source_item": idx   # align with dataset via index
-            })
-
-    # --- Priority Rule: Dataset spans win over relation overlaps ---
-    def overlaps(a, b):
-        return not (a["end"] <= b["start"] or a["start"] >= b["end"])
-
-    dataset_spans = sorted(dataset_spans, key=lambda s: s["start"])
-
-    filtered_relations = []
-    for rel in relation_spans:
-        if any(overlaps(rel, ds) for ds in dataset_spans):
-            continue
-        filtered_relations.append(rel)
-
-    # --- Deduplicate ---
-    uniq = {(s["start"], s["end"], s["label"], s["kind"]): s
-            for s in dataset_spans + filtered_relations}
-
-    spans = sorted(uniq.values(), key=lambda s: s["start"])
-    return spans
-
-
-def render_highlight(text):
-    res = extract_from_text(text, use_classifier_gate=False)
-    base = res[0]["text"]
-
-    spans = collect_spans(res)
-
-    if not spans:
-        return f"<div style='white-space:pre-wrap;'>{base}</div>"
-
-    html = """<style>
-    .hl {
-    position: relative;
-    cursor: pointer;
-    display: inline-block;
-    }
-
-    .hl:hover::after {
-    content: attr(data-tip);
-    position: absolute;
-    top: -2.2em;
-    left: 0;
-    background: #fff;           /* white tooltip */
-    color: #222;                /* dark text */
-    padding: 4px 8px;
-    font-size: 0.78em;
-    border-radius: 6px;
-    border: 1px solid #ccc;     /* subtle border */
-    box-shadow: 0px 2px 6px rgba(0,0,0,0.15);  /* soft shadow */
-    white-space: nowrap;
-    z-index: 9999;
-    opacity: 1;
-    }
-
-    .hl::after {
-    opacity: 0;
-    pointer-events: none;
-    transition: opacity 0.05s; /* still instant, but fades cleanly */
-    }
-    </style>
-    <div style='white-space:pre-wrap; font-family:monospace;'>
+            for t in sorted(targets):
+                html_out.append(
+                    f"<li><span style='color:#80c904;'>{html.escape(rel_type)}</span>: {html.escape(t)}</li>"
+                )
+        html_out.append("</ul></details>")
+    html_out.append("</div>")
+
+    return "\n".join(html_out)
+
+import html
+from collections import defaultdict
+# ========================================
+# Build Tree from Deduplicated Mentions
+# ========================================
+import html
+import re
+
+def make_relation_tree_from_dedup(dedup_results):
     """
+    Build a cleaned, deduplicated relation tree with prefiltering.
+    - Publishers that are dataset names are removed
+    - Years validated with regex \b\d{4}\b
+    - Empty/noisy values removed
+    """
+
+    if not dedup_results or not isinstance(dedup_results, list):
+        return "<i>No deduplicated results found.</i>"
+
+    # --- prepare dataset name list for filtering publisher noise ---
+    dataset_names = {d.get("text", "").lower() for d in dedup_results if d.get("text")}
+    relation_priority = [
+        "acronym", "author", "publisher", "data type", "data description",
+        "reference year", "publication year", "data geography",
+        "reference population", "usage context"
+    ]
+
+    def _filter_values(rel_type, values):
+        """relation-specific filtering logic"""
+        if not values:
+            return []
+
+        # ensure list
+        if isinstance(values, str):
+            values = [values]
+        elif not isinstance(values, (list, set, tuple)):
+            return []
+
+        clean = []
+        for v in values:
+            if not v or not isinstance(v, str):
+                continue
+            val = v.strip()
+            if not val:
+                continue
+
+            # publisher should not be dataset name
+            if rel_type == "publisher":
+                if val.lower() in dataset_names:
+                    continue
+                # also skip if it's a dataset acronym
+                if any(val.lower() in str(a).lower() for d in dedup_results for a in d.get("acronym", [])):
+                    continue
+
+            # year fields must contain 4-digit year
+            if rel_type in {"reference year", "publication year"}:
+                if not re.search(r"\b\d{4}\b", val):
+                    continue
+
+            # simple length sanity check
+            if len(val) > 150:
+                continue
 
+            clean.append(val)
+        return sorted(set(clean))
 
-    last = 0
-
-    for span in spans:
-        html += base[last:span["start"]]
-
-        if span["kind"] == "dataset":
-            color = DATASET_COLORS.get(span["label"], "#999")
-        else:
-            color = RELATION_COLORS.get(span["label"], "#74c0fc")
-
-        # if span["kind"] == "dataset":
-        #     tooltip = span["label"]
-        if span["kind"] == "dataset":
-            tooltip = span["label"] if not span.get("usage") else f"usage context: {span['usage']}"
-
-        else:
-            # Show which dataset this relation corresponds to
-            # We look up the nearest dataset with same extraction item.
-            # But since source is provided already, use it directly:
-            tooltip = f"{span['label']} for {span['source']}"
-
-        display_label = RELATION_DISPLAY.get(span["label"], span["label"])
-        html += (
-            f"<span class='hl' data-tip='{tooltip}' "
-            f"style='background:{color}22; border:1px solid {color}; "
-            f"border-radius:4px; padding:0px 3px; margin:0 2px;'>"
-            f"{span['text']}"
-            f"<span style='border:1px solid {color}; color:{color}; "
-            f"font-size:0.7em; border-radius:3px; padding:1px 3px; margin-left:5px;'>"
-            f"{display_label}</span></span>"
+    # --- Build HTML output ---
+    html_out = ["<div style='font-family:monospace; font-size:0.9em; line-height:1.5;'>"]
+
+    for entry in dedup_results:
+        ds_name = entry.get("text", "").strip()
+        ds_label = entry.get("label", "")
+        if not ds_name:
+            continue
+
+        ds_badge = f" <span style='color:gray;'>[{html.escape(ds_label)}]</span>" if ds_label else ""
+        html_out.append(
+            f"<details open><summary><b style='color:#4da6ff;'>{html.escape(ds_name)}</b>{ds_badge}</summary><ul>"
         )
-        last = span["end"]
-
-    html += base[last:]
-    html += "</div>"
-
-    return html
-
-with gr.Blocks(
-    title="AI for Data Use: Dataset Extraction",
-    css="""
-    body, .gradio-container { background-color: #111 !important; color: #e6e6e6 !important; }
-
-    textarea, input, .gr-textbox {
-        background: #1a1a1a !important;
-        color: #f2f2f2 !important;
-        border: 1px solid #444 !important;
-        border-radius: 6px !important;
-    }
-    textarea:focus, input:focus {
-        border-color: #888 !important;
-        box-shadow: 0 0 0 1px #888 !important;
-    }
-
-    button {
-        background: #333 !important;
-        color: #eee !important;
-        border-radius: 6px !important;
-        border: 1px solid #444 !important;
-        padding: 8px 14px !important;
-    }
-    button:hover {
-        background: #444 !important;
-        border-color: #666 !important;
-    }
-
-    .gr-html {
-        background: #1a1a1a !important;
-        color: #fff !important;
-        border-radius: 6px !important;
-        border: 1px solid #333 !important;
-        padding: 14px !important;
-        min-height: 200px;
-    }
-    """
-) as demo:
 
-    gr.Markdown("# AI for Data Use: Dataset Extraction")
+        rel_keys = [k for k in entry.keys() if k not in {
+            "text", "label", "score", "count", "form_counts", "other_datasets",
+            "mentioned_in_list", "mentioned_in_sentence_list", "pages",
+            "sources", "start_indices", "end_indices", "raw_contexts"
+        }]
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            input_box = gr.Textbox(lines=14, label="Input Text")
+        rel_keys = sorted(
+            rel_keys,
+            key=lambda r: relation_priority.index(r) if r in relation_priority else len(relation_priority)
+        )
 
-        with gr.Column(scale=1):
-            output_box = gr.HTML(
-                label="Highlighted Output",
-                value="<div style='min-height:200px; border:1px solid #444; border-radius:8px; padding:10px; opacity:0.6; text-align:center;'>Waiting for input…</div>"
+        for rel_key in rel_keys:
+            vals = _filter_values(rel_key, entry.get(rel_key))
+            if not vals:
+                continue
+            html_out.append(
+                f"<li><span style='color:#80c904;'>{html.escape(rel_key)}</span>: "
+                f"{html.escape('; '.join(vals))}</li>"
             )
 
-    submit_btn = gr.Button("Submit")
+        html_out.append("</ul></details>")
 
-    submit_btn.click(
-        fn=render_highlight,
-        inputs=input_box,
-        outputs=output_box
+    html_out.append("</div>")
+    return "\n".join(html_out)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧠 AI for Data Use: Dataset Extraction (Deduplicated Tree)")
+
+    inp = gr.Textbox(
+        label="Input Text",
+        lines=10,
+        placeholder="Paste or type a paragraph here..."
     )
 
+    run_btn = gr.Button("🚀 Run Extraction")
+    out = gr.HighlightedText(label="Highlighted Datasets and Relations")
+
+    tree_btn = gr.Button("🧭 Show / Refresh Relation Tree")
+    tree_html = gr.HTML(label="Relation Tree", visible=False)
+
+    model_state = gr.State()      # raw extract_from_text() results
+    highlight_state = gr.State()  # processed highlights
+
     gr.Examples(
-        examples=[
-    ["We examine early childhood nutrition and health inequality across Sub-Saharan Africa by combining information from the Demographic and Health Surveys (DHS, various years) and the Multiple Indicator Cluster Surveys (MICS). Anthropometric outcomes (height-for-age, weight-for-age) are standardized using WHO Child Growth Standards. To account for environmental exposure, we align survey cluster coordinates with gridded temperature and precipitation data from the CRU TS 4.06 dataset at the nearest spatial cell."],
-    ["Introduction The mining sector in Africa is growing rapidly and is the main recipient of foreign direct investment ( World Bank 2011 ). The welfare effects of this sector are not well understood, although a literature has recently developed around this question. The main contribution of this paper is to shed light on the welfare effects of gold mining in a detailed, in-depth country study of Ghana, a country with a long tradition of gold mining and a recent, large expansion in capital - intensive and industrial-scale production. A second contribution of this paper is to show the importance of decomposing the effects with respect to distance from the mines. Given the spatial heterogeneity of the results, we explore the effects in an individual-level, difference-in-differences analysis by using spatial lag models to allow for nonlinear effects with distance from mine. We also allow for spillovers across districts, in a district-level analysis. We use two complementary geocoded household data sets to analyze outcomes in Ghana: the Demographic and Health Survey ( DHS ) and the Ghana Living Standard Survey ( GLSS ), which provide information on a wide range of welfare outcomes."],
-    ["The main mining data is a dataset from InterraRMG covering all large-scale mines in Ghana, explained in more detail in section 3. 1. This dataset is linked to survey data from the DHS and GLSS, using spatial information. Geographical coordinates of enumeration areas in GLSS are from Ghana Statistical Services ( GSS ). 2 Point coordinates ( global positioning system [ GPS ] ) for the surveyed DHS clusters3 allow us to match all individuals to one or several mineral mines. We do this in two ways. First, we calculate distance spans from an exact mine location given by its GPS coordinates, and match surveyed individuals to mines. "],
-    ["We study learning outcomes by linking standardized test score data from the Programme for International Student Assessment (PISA) and the Trends in International Mathematics and Science Study (TIMSS) to school resource indicators compiled by the UNESCO Institute for Statistics (UIS). Household socioeconomic background is proxied using parental education information extracted from the Demographic and Health Surveys (DHS) female and household datasets. To examine the role of inequality, we calculate within-school and between-school variance components and correlate these with school funding data from national Ministry of Education financial reports. The goal is to understand how unequal resource allocation contributes to learning gaps across income groups and geographic regions."],
-    ["Patterns of forced displacement are analyzed using the UNHCR Refugee Population Statistics Database, which provides annual country-to-country asylum flows and refugee stock counts. To understand displacement drivers, we integrate conflict event data from the ACLED conflict monitoring system and national political stability indices from the World Governance Indicators (WGI). We additionally incorporate bilateral migration dyads from the World Bank Global Bilateral Migration Database to capture historical migration ties. Geographic exposure to violence is assigned using region-level coordinates and spatial joining procedures. This dataset allows us to estimate the relationship between escalating conflict intensity and subsequent cross-border population movements."],
-    ],
-        inputs=input_box
+        examples=example_texts,
+        inputs=inp,
+        label="🧪 Try these examples"
     )
 
-    gr.Markdown("""
+    # --- main pipeline ---
+    def run_pipeline(text):
+        text = truncate_text(text, max_tokens=300)
+        highlights, results = process_text(text)
+        return highlights, results, highlights
+
+    run_btn.click(
+        fn=run_pipeline,
+        inputs=inp,
+        outputs=[out, model_state, highlight_state]
+    ).then(
+        fn=lambda h, m: gr.update(
+            visible=True,
+            value=make_relation_tree_from_dedup(deduplicate_mentions(m))
+        ),
+        inputs=[highlight_state, model_state],
+        outputs=tree_html
+    )
+
+    # --- manual refresh ---
+    def refresh_relation_tree(h, m):
+        if not h or not m:
+            return gr.update(visible=True, value="<i>No extracted data yet — run extraction first.</i>")
+        html_tree = make_relation_tree_from_dedup(deduplicate_mentions(m))
+        return gr.update(visible=True, value=html_tree)
+
+    tree_btn.click(fn=refresh_relation_tree, inputs=[highlight_state, model_state], outputs=tree_html)
+
+    gr.Markdown(f"""
 <hr style='border: none; border-top: 1px solid #333; margin: 20px 0;'>
 
 ### AI for Data Use: Dataset Extraction
@@ -358,6 +537,5 @@ The detected spans are **highlighted inline** with labels so you can quickly rev
 - Project Docs: https://worldbank.github.io/ai4data-use/docs/introduction.html
 """)
 
-demo.launch()
-
 
+demo.launch()

requirements.txt

@@ -1,7 +1 @@
-gradio>=4.0.0
-datasets
-transformers
-accelerate
-sentencepiece
-gliner
-huggingface_hub
+gradio>=4.0.0