Spaces:
Sleeping
Sleeping
| import re | |
| import gradio as gr | |
| from collections import defaultdict | |
| import html | |
| from typing import List | |
| import os | |
| import subprocess | |
| import sys | |
| # # Ensure correct hub version before anything else | |
| # subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "huggingface_hub==0.24.5"], check=True) | |
| # # Optional: safety patch | |
| # import huggingface_hub.constants as constants | |
| # if not hasattr(constants, "HF_HUB_ENABLE_HF_TRANSFER"): | |
| # constants.HF_HUB_ENABLE_HF_TRANSFER = False | |
| GH_TOKEN = os.getenv("GH_TOKEN") | |
| if GH_TOKEN: | |
| subprocess.run( | |
| [ | |
| "pip", "install", | |
| f"git+https://{GH_TOKEN}:x-oauth-basic@github.com/rafmacalaba/ai4data_use.git" | |
| ], | |
| check=True | |
| ) | |
| else: | |
| print("GH_TOKEN not found. Private ai4data_use will NOT be installed.") | |
| from ai4data import extract_from_text, deduplicate_mentions | |
| DATA_MODEL_ID = "rafmacalaba/datause-extraction-v3-finetuned" | |
| 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: | |
| 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 | |
| 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 | |
| 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 = {} | |
| for item in results: | |
| ds = item.get("datasets") | |
| rels = item.get("relations", []) | |
| if not ds: | |
| continue | |
| ds_name = ds.get("text", "").strip() | |
| ds_label = ds.get("label", "") | |
| if not ds_name: | |
| continue | |
| # Initialize dataset node | |
| if ds_name not in tree: | |
| tree[ds_name] = {"_dataset_label": ds_label} | |
| # Collect all relations | |
| for rel in rels: | |
| rel_type = rel.get("relation", "").strip().lower() | |
| target = rel.get("target", "").strip() | |
| # Skip empty or malformed relations | |
| if not rel_type or not target: | |
| continue | |
| # Special handling for usage context (value-only) | |
| if rel_type == "usage context" and "value" in rel: | |
| target = rel["value"] | |
| # Attach relation under this dataset | |
| tree[ds_name].setdefault(rel_type, set()).add(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): | |
| 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)) | |
| # --- 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>" | |
| ) | |
| 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" | |
| }] | |
| rel_keys = sorted( | |
| rel_keys, | |
| key=lambda r: relation_priority.index(r) if r in relation_priority else len(relation_priority) | |
| ) | |
| 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>" | |
| ) | |
| html_out.append("</ul></details>") | |
| html_out.append("</div>") | |
| return "\n".join(html_out) | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## 🧠 AI for Data Use: Dataset Extraction") | |
| 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=example_texts, | |
| inputs=inp, | |
| label="🧪 Try these examples" | |
| ) | |
| # --- 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 | |
| This tool identifies **dataset mentions** (e.g., *Demographic and Health Survey*, *Living Standards and Measurement Survey*, etc.) and extracts **contextual metadata** such as: | |
| - publisher | |
| - publication year | |
| - reference year | |
| - geography | |
| - acronym | |
| - reference population | |
| - data description | |
| - data type | |
| - usage context | |
| **Usage Context Definitions** | |
| - **Primary mention** – the dataset is the **main source of analysis or results** in the study. | |
| - **Supporting mention** – the dataset is **used alongside other data** to complement or validate findings. | |
| - **Background mention** – the dataset is **mentioned for context or comparison** but **not used** in the actual analysis. | |
| #### How to Use | |
| 1. Paste or type text into the input box (left), or select one of the provided examples. | |
| 2. Click **🚀 Run Extraction** to process the text. | |
| 3. The model will **highlight all detected dataset mentions** and related entities (e.g., publisher, geography, year, usage context) directly in the text. | |
| 4. Below the highlights, a **deduplicated relation tree** will automatically appear, showing each dataset with its extracted metadata and filtered attributes. | |
| 5. You can click **🧭 Show / Refresh Relation Tree** anytime to rebuild or inspect the deduplicated metadata view. | |
| <hr style='border: none; border-top: 1px solid #333; margin: 20px 0;'> | |
| **Resources** | |
| - Model: https://huggingface.co/{DATA_MODEL_ID} | |
| - Paper (ArXiv): https://arxiv.org/pdf/2502.10263 | |
| - GLiNER Repo: https://github.com/urchade/GLiNER | |
| - Project Docs: https://worldbank.github.io/ai4data-use/docs/introduction.html | |
| """) | |
| demo.launch() |