Update app.py

app.py

@@ -1,174 +1,335 @@
+"""Pl@ntBERT's Gradio interface. Run `python app.py` for a local preview."""
+
+import atexit
+import csv
+import json
+import logging
+import os
+import shutil
+import tempfile
+import time
+from pathlib import Path
+
+os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
+os.environ.setdefault("HF_HUB_OFFLINE", "1")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
 import gradio as gr
-from transformers import pipeline
-import requests
-import pygbif
-from bs4 import BeautifulSoup
-import pandas as pd
-
-classification_model = pipeline("text-classification", model="models/text_classification_model", tokenizer="models/text_classification_model", top_k=5)
-masking_model = pipeline("fill-mask", model="models/fill_mask_model", tokenizer="models/fill_mask_model", top_k=100)
-
-eunis_habitats = pd.read_excel('data/eunis_habitats.xlsx')
-
-image_not_found = gr.Image("https://img.freepik.com/premium-vector/file-folder-mascot-character-design-vector_166742-4413.jpg")
-
-def return_image(task, label):
-    image = None
-    if task == "classification":
-        floraveg_url = f"https://floraveg.eu/habitat/overview/{label}"
-        floraveg_tag = "https://files.ibot.cas.cz/cevs/images/syntaxa/thumbs/"
-    elif task == "masking":
-        floraveg_url = f"https://floraveg.eu/taxon/overview/{label}"
-        floraveg_tag = "https://files.ibot.cas.cz/cevs/images/taxa/large/"
-    response = requests.get(floraveg_url)
-    if response.status_code == 200:
-        soup = BeautifulSoup(response.text, 'html.parser')
-        img_tag = soup.find('img', src=lambda x: x and x.startswith(floraveg_tag))
-        if img_tag:
-            image_url = img_tag['src']
-            image = gr.Image(value=image_url)
-    return image
-
-def gbif_normalization(text):
-    all_species = text.split(',')
-    all_species = [species.strip() for species in all_species]
-    species_gbif = []
-    for species in all_species:
-        gbif_match_result = pygbif.species.name_backbone(species, taxonRank="SPECIES")
-        if 'usage' in gbif_match_result:
-            r = gbif_match_result["usage"]["canonicalName"]
-        else:
-            r = species
-        species_gbif.append(r)
-    text = ", ".join(species_gbif)
-    text = text.lower()
-    return text
-
-def classification(text, k):
-    text = gbif_normalization(text)
-    result = classification_model(text)
-    habitat_labels = [res['label'] for res in result[0][:k]]
-    habitat_name = eunis_habitats[eunis_habitats['EUNIS 2020 code'] == habitat_labels[0]]['EUNIS-2021 habitat name'].values[0]
-    if k == 1:
-        text = f"This vegetation plot probably belongs to the habitat type {habitat_labels[0]}."
-        text += f"\nThis habitat type is named '{habitat_name}'."
-    elif k == 2:
-        text = f"This vegetation plot probably belongs to the habitat type {', '.join(habitat_labels[:-1])} or {habitat_labels[-1]}."
-        text += f"\nThe most likely habitat type (i.e., {habitat_labels[0]}) is named '{habitat_name}'."
-    else:
-        text = f"This vegetation plot probably belongs to the habitat type {', '.join(habitat_labels[:-1])}, or {habitat_labels[-1]}."
-        text += f"\nThe most likely habitat type (i.e., {habitat_labels[0]}) is named '{habitat_name}'."
-    image = return_image("classification", habitat_labels[0])
-    if image is not None:
-        text += "\nBelow is an example of this habitat type taken from the website FloraVEG."
-    else:
-        text += f"\nNo image found for this habitat type."
-        image = image_not_found
-    return text, image
-
-def masking(text, k):
-    text = gbif_normalization(text)
-    text_split = text.split(', ')
-    
-    best_predictions = []
-    
-    for _ in range(k):
-        max_score = 0
-        best_prediction = None
-        best_position = None
-        best_sentence = None
-
-        for i in range(len(text_split) + 1):
-            masked_text = ', '.join(text_split[:i] + ['[MASK]'] + text_split[i:])
-            
-            j = 0
-            while True:
-                prediction = masking_model(masked_text)[j]
-                species = prediction['token_str']
-                if species in text_split or species in best_predictions:
-                    j += 1
-                else:
-                    break
-
-            score = prediction['score']
-            sentence = prediction['sequence']
-
-            if score > max_score:
-                max_score = score
-                best_prediction = species
-                best_position = i
-                best_sentence = sentence
-        
-        best_predictions.append(best_prediction)
-        text_split.insert(best_position, best_prediction)
-        
-    best_positions = [text_split.index(prediction) for prediction in best_predictions]
-    best_predictions = [s.strip().capitalize() for s in best_predictions]
-    best_sentence = ", ".join([s.strip().capitalize() for s in best_sentence.split(",")])
-    
-    if k == 1:
-        text = f"The most likely missing species is '{best_predictions[0]}' (position {best_positions[0]})."
-    elif k == 2:
-        text = f"The most likely missing species are '{best_predictions[0]}' and '{best_predictions[1]}' (positions {best_positions[0]} and {best_positions[1]})."
-    else:
-        text = f"The most likely missing species are " + ', '.join(f"'{s}'" for s in best_predictions[:-1]) + f", and '{best_predictions[-1]}' (positions {', '.join(map(str, best_positions[:-1]))}, and {best_positions[-1]})."
-    text += f"\nThe completed vegetation plot is thus '{best_sentence}'."
-    image = return_image("masking", best_predictions[0])
-    if image is not None:
-        text += f"\nBelow is an image of the first missing species (i.e., {best_predictions[0]}) taken from the website FloraVEG."
-    else:
-        text += f"\nNo image found for the first missing species (i.e., {best_predictions[0]})."
-        image = image_not_found
-    return text, image
-
-with gr.Blocks() as demo:
-
-    gr.Markdown("""<h1 style="text-align: center;">Pl@ntBERT</h1>""")
-    
-    with gr.Tab("Vegetation plot classification"):
-        gr.Markdown("""<h3 style="text-align: center;">Habitat identification of vegetation plots!</h3>""")
-        with gr.Row():
-            with gr.Column():
-                species_classification = gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here.")
-                k_classification = gr.Slider(1, 5, value=1, step=1, label="Top-k", info="Choose the number of habitat types to display.")
-            with gr.Column():
-                text_classification = gr.Textbox(label="Prediction")
-                image_classification = gr.Image()
-        button_classification = gr.Button("Classify")
-        gr.Markdown("""<h5 style="text-align: center;">An example of input</h5>""")
-        gr.Examples(
-            examples=[
-                ["phragmites australis, lemna minor, typha latifolia", 3],
-                ["fagus sylvatica, prenanthes purpurea, abies alba, cardamine bulbifera, orthilia secunda, oxalis acetosella, rubus idaeus", 1]
-            ], 
-            inputs=[species_classification, k_classification],
-            outputs=[text_classification, image_classification],
-            fn=classification, 
-            cache_examples=True)
-        
-    with gr.Tab("Missing species finding"):
-        gr.Markdown("""<h3 style="text-align: center;">Missing vascular plant species retrieval!</h3>""")
-        with gr.Row():
-            with gr.Column():
-                species_masking = gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here.")
-                k_masking = gr.Slider(1, 5, value=1, step=1, label="Top-k", info="Choose the number of missing species to find.")
-            with gr.Column():
-                text_masking = gr.Textbox(label="Prediction")
-                image_masking = gr.Image()
-        button_masking = gr.Button("Find")
-        gr.Markdown("""<h5 style="text-align: center;">An example of input</h5>""")
-        gr.Examples(
-            examples=[
-            ["calamagrostis arenaria, medicago marina, pancratium maritimum, thinopyrum junceum", 1],
-            ["trapa natans, lemna minor, phragmites australis, sparganium erectum", 1]
-            ],
-            inputs=[species_masking, k_masking],
-            outputs=[text_masking, image_masking], 
-            fn=masking, 
-            cache_examples=True
+
+from illustrations import render_illustration
+from plantbert import ROOT, InputError, ModelUnavailable, TaxonomyError, run_inference
+from presentation import (
+    ABOUT,
+    ACKNOWLEDGEMENTS,
+    CONTACT,
+    HERO,
+    empty_result,
+    error_result,
+    render_audit,
+    render_normalization,
+    render_results,
+)
+
+LOGGER = logging.getLogger(__name__)
+EXPORT_DIRECTORY = Path(tempfile.mkdtemp(prefix="plantbert-exports-"))
+atexit.register(shutil.rmtree, EXPORT_DIRECTORY, ignore_errors=True)
+
+EXAMPLES = {
+    "habitat": [
+        ("Tall-helophyte bed", "Phragmites australis, Lemna minor, Typha latifolia"),
+        (
+            "Fagus forest on non-acid soils",
+            "Fagus sylvatica, Prenanthes purpurea, Abies alba, Cardamine bulbifera, Orthilia secunda, Oxalis acetosella, Rubus idaeus",
+        ),
+    ],
+    "species": [
+        (
+            "Eryngium maritimum",
+            "Calamagrostis arenaria, Medicago marina, Pancratium maritimum, Thinopyrum junceum",
+        ),
+        (
+            "Typha angustifolia",
+            "Trapa natans, Lemna minor, Phragmites australis, Sparganium erectum",
+        ),
+    ],
+}
+
+
+def export_result(result: dict) -> str:
+    for old in EXPORT_DIRECTORY.glob("*"):
+        if old.stat().st_mtime < time.time() - 3600:
+            old.unlink(missing_ok=True)
+    with tempfile.NamedTemporaryFile(
+        mode="w",
+        suffix=".json",
+        prefix=f"plantbert-{result['task']}-",
+        dir=EXPORT_DIRECTORY,
+        delete=False,
+        encoding="utf-8",
+    ) as handle:
+        json.dump(result, handle, ensure_ascii=False, indent=2, allow_nan=False)
+        return handle.name
+
+
+def export_csv(result):
+    with tempfile.NamedTemporaryFile(
+        mode="w",
+        suffix=".csv",
+        prefix="plantbert-results-",
+        dir=EXPORT_DIRECTORY,
+        delete=False,
+        encoding="utf-8",
+        newline="",
+    ) as handle:
+        fields = [
+            "rank_or_step",
+            "code",
+            "name",
+            "model_score",
+            "final_position",
+            "normalized_input",
+        ]
+        writer = csv.DictWriter(handle, fieldnames=fields)
+        writer.writeheader()
+        for rank, prediction in enumerate(result["predictions"], 1):
+            writer.writerow(
+                {
+                    "rank_or_step": prediction.get("step", rank),
+                    "code": prediction.get("code", ""),
+                    "name": prediction.get("name", prediction.get("species", "")),
+                    "model_score": prediction["score"],
+                    "final_position": prediction.get("final_position", ""),
+                    "normalized_input": ", ".join(result["input"]["names"]),
+                }
+            )
+        return handle.name
+
+
+def analyze(text, k, task, progress):
+    try:
+        result = run_inference(text, k, task, progress)
+        return (
+            render_results(result),
+            render_audit(result),
+            result,
+            gr.update(value=export_result(result), interactive=True),
+            gr.update(value=export_csv(result), interactive=True),
+        )
+    except (InputError, ModelUnavailable) as error:
+        audit = render_normalization(error.entries) if isinstance(error, TaxonomyError) else ""
+        return (
+            error_result(str(error)),
+            audit,
+            None,
+            gr.update(value=None, interactive=False),
+            gr.update(value=None, interactive=False),
+        )
+    except Exception:
+        LOGGER.exception("Pl@ntBERT inference failed")
+        return (
+            error_result(
+                "The model could not complete this request. Please retry with a shorter plot. "
+                "If the problem persists, contact the developer through Methods & citation."
+            ),
+            "",
+            None,
+            gr.update(value=None, interactive=False),
+            gr.update(value=None, interactive=False),
+        )
+
+
+def classify(text, k, progress=gr.Progress()):  # noqa: B008 - Gradio injects this default
+    return analyze(text, k, "habitat", progress)
+
+
+def complete(text, k, progress=gr.Progress()):  # noqa: B008 - Gradio injects this default
+    return analyze(text, k, "species", progress)
+
+
+def clear_results(task):
+    return (
+        empty_result(task),
+        "Run a prediction to review the names used by the model.",
+        None,
+        gr.update(value=None, interactive=False),
+        gr.update(value=None, interactive=False),
+        "",
+    )
+
+
+def build_task(task):
+    habitat = task == "habitat"
+    with gr.Row(equal_height=True, elem_classes="workspace"):
+        with gr.Column(scale=1, min_width=300, elem_classes=["work-panel", "input-panel"]):
+            gr.HTML('<div class="section-kicker">01 &nbsp; YOUR VEGETATION PLOT</div>')
+            gr.Markdown("### Which plants occur together?")
+            species = gr.Textbox(
+                label="Scientific plant names",
+                lines=4,
+                max_lines=10,
+                placeholder="Phragmites australis, Lemna minor, Typha latifolia",
+                value=EXAMPLES[task][0][1],
+                info="Separate names with commas or new lines. Order by decreasing cover when known.",
+                elem_classes="species-input",
+            )
+            gr.HTML('<p class="example-label">START WITH AN EXAMPLE</p>')
+            example_buttons = []
+            with gr.Row(elem_classes="example-buttons"):
+                for label, text in EXAMPLES[task]:
+                    button = gr.Button(label, size="sm", variant="secondary")
+                    example_buttons.append(button)
+                    button.click(
+                        lambda value=text: value,
+                        outputs=species,
+                        queue=False,
+                        api_visibility="private",
+                    )
+            k = gr.Slider(
+                1,
+                5,
+                value=3 if habitat else 1,
+                step=1,
+                label="Habitats to display" if habitat else "Species to add sequentially",
+                info="Compare the highest-scoring habitat classes."
+                if habitat
+                else "Each addition informs the next. More additions take longer.",
+            )
+            run = gr.Button(
+                "Identify habitat  →" if habitat else "Suggest missing species  →",
+                variant="primary",
+                elem_classes="run-button",
+            )
+        with gr.Column(scale=1, min_width=300, elem_classes=["work-panel", "output-panel"]):
+            gr.HTML('<div class="section-kicker">02 &nbsp; MODEL EXPLORATION</div>')
+            result = gr.HTML(empty_result(task))
+            photo = gr.HTML()
+    with gr.Column(elem_classes="record-panel"):
+        gr.HTML(
+            '<div class="section-kicker">03 &nbsp; CHECK NAMES & SAVE RESULTS</div>'
+            '<p class="record-intro">Review the species used, then take your results with you.</p>'
         )
+        audit = gr.HTML("Run a prediction to review the names used by the model.")
+        with gr.Row(elem_classes="download-row"):
+            csv_download = gr.DownloadButton(
+                "Download results · CSV", interactive=False, variant="primary"
+            )
+            download = gr.DownloadButton("Download full record · JSON", interactive=False)
+        with gr.Accordion(
+            "Scientific record · settings, name checks & checkpoint provenance", open=False
+        ):
+            record = gr.JSON(label="Reproducibility record", value=None)
+    outputs = [result, audit, record, download, csv_download]
+    controls = [species, k, run, *example_buttons]
+    run.click(
+        lambda: [gr.update(interactive=False) for _ in controls],
+        outputs=controls,
+        queue=False,
+        api_visibility="private",
+    ).then(
+        classify if habitat else complete,
+        inputs=[species, k],
+        outputs=outputs,
+        api_name="classify_habitat" if habitat else "suggest_species",
+        concurrency_limit=1,
+        concurrency_id="plantbert-cpu",
+        trigger_mode="once",
+    ).then(
+        render_illustration,
+        inputs=record,
+        outputs=photo,
+        concurrency_id="illustrations",
+        concurrency_limit=4,
+        api_visibility="private",
+    ).then(
+        lambda: [gr.update(interactive=True) for _ in controls],
+        outputs=controls,
+        queue=False,
+        api_visibility="private",
+    )
+    for component in (species, k):
+        component.change(
+            lambda: clear_results(task),
+            outputs=[*outputs, photo],
+            queue=False,
+            api_visibility="private",
+        )
+
+
+def build_demo():
+    with gr.Blocks(
+        title="Pl@ntBERT · Plant community explorer",
+        analytics_enabled=False,
+        delete_cache=(3600, 3600),
+    ) as demo:
+        with gr.Row(elem_classes="theme-toolbar"):
+            theme_toggle = gr.Button("Dark mode", size="sm", elem_id="theme-toggle")
+        demo.load(fn=None, js=(ROOT / "assets/theme.js").read_text(), queue=False)
+        theme_toggle.click(
+            fn=None,
+            js=(ROOT / "assets/theme-change.js").read_text(),
+            queue=False,
+        )
+        gr.HTML(HERO)
+        with gr.Tabs(elem_id="task-tabs"):
+            with gr.Tab("Identify a habitat", id="habitat"):
+                build_task("habitat")
+            with gr.Tab("Find missing species", id="species"):
+                build_task("species")
+            with gr.Tab("Methods & citation", id="methods"):
+                gr.Markdown(ABOUT, elem_classes="methods")
+                gr.HTML(ACKNOWLEDGEMENTS, elem_classes="methods-funding")
+                with gr.Accordion(
+                    "Contact the developer", open=False, elem_classes="contact-accordion"
+                ):
+                    gr.HTML(CONTACT)
+        gr.HTML(
+            '<footer class="research-footer"><span><strong>Pl@ntBERT</strong> &nbsp; · &nbsp; Learning the syntax of plant assemblages</span>'
+            "<span>Research demo &nbsp; / &nbsp; Interpret with ecological context</span></footer>"
+        )
+    return demo.queue(max_size=12, default_concurrency_limit=1)
+
+
+THEME = gr.themes.Base(
+    primary_hue="emerald",
+    secondary_hue="stone",
+    neutral_hue="stone",
+    font=["Inter", "ui-sans-serif", "system-ui", "sans-serif"],
+    font_mono=["ui-monospace", "monospace"],
+    radius_size="lg",
+).set(
+    body_background_fill="#f8f9f5",
+    body_background_fill_dark="#14241f",
+    body_text_color="#203c32",
+    body_text_color_dark="#e2ede6",
+    block_background_fill="#ffffff",
+    block_background_fill_dark="#1c3028",
+    block_border_color="#dce5dd",
+    block_border_color_dark="#3c5146",
+    button_primary_background_fill="#24664d",
+    button_primary_background_fill_hover="#184b38",
+    button_primary_text_color="#ffffff",
+    button_primary_border_color="#24664d",
+    button_primary_background_fill_dark="#79b99a",
+    button_primary_text_color_dark="#10271d",
+    button_secondary_background_fill="#f1f5ee",
+    button_secondary_background_fill_dark="#263e32",
+    input_background_fill="#fbfcf9",
+    input_background_fill_dark="#14271e",
+    input_border_color="#cfdbd0",
+    input_border_color_focus="#327557",
+    block_shadow="none",
+    block_title_text_weight="500",
+)
 
-    button_classification.click(classification, inputs=[species_classification, k_classification], outputs=[text_classification, image_classification])
-    button_masking.click(masking, inputs=[species_masking, k_masking], outputs=[text_masking, image_masking])
 
-demo.launch()
+if __name__ == "__main__":
+    build_demo().launch(
+        theme=THEME,
+        css_paths=[ROOT / "assets" / "style.css"],
+        server_name="0.0.0.0" if os.environ.get("SPACE_ID") else "127.0.0.1",
+        server_port=int(os.environ.get("PORT", "7860")),
+        share=False,
+        ssr_mode=False,
+        footer_links=[],
+        show_error=False,
+    )