app.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / CUDA-touching imports
import re
import torch
import gradio as gr
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

BASE_MODEL = "Qwen/Qwen3.5-2B-Base"
ADAPTER = "EdisonScientific/MarkushGlyph"
MAX_NEW_TOKENS_DEFAULT = 1024
IMAGE_MAX_PIXELS = 1_048_576  # 1024x1024 — matches glyph checkpoint loading

ENDOFTEXT_TOKEN_ID = 248044
IM_END_TOKEN_ID = 248046

PROMPT = (
    "Extract the Markush structure from this image.\n"
    "Target format: return <markush><cxsmi>...</cxsmi><stable>...</stable></markush>; "
    "the <cxsmi> value must use canonicalized cxsmiles_opt."
)

# ---------------------------------------------------------------------------
# Prediction cleaning (ported from glyph.markush.eval.common)
# ---------------------------------------------------------------------------

_XML_FLAGS = re.DOTALL | re.IGNORECASE

RE_THINKING = re.compile(r"<think>.*?</think>\s*", _XML_FLAGS)
RE_CXSMI = re.compile(r"<cxsmi>(.*?)</cxsmi>", _XML_FLAGS)
RE_STABLE = re.compile(r"<stable>(.*?)</stable>", _XML_FLAGS)
RE_MARKUSH = re.compile(
    r"<markush>\s*<cxsmi>(.*?)</cxsmi>\s*(?:<stable>(.*?)</stable>)?\s*</markush>",
    _XML_FLAGS,
)

_CHAT_TEMPLATE_TOKENS = ("<|im_end|>", "<|endoftext|>", "<|im_start|>")


def clean_prediction(text: str | None) -> str:
    """Strip thinking blocks and chat-template artifacts from raw model output."""
    if not text:
        return ""
    text = RE_THINKING.sub("", text)
    for tok in _CHAT_TEMPLATE_TOKENS:
        text = text.split(tok)[0]
    return text.strip()


# ---------------------------------------------------------------------------
# Processor helpers (ported from glyph.markush.eval.checkpoint)
# ---------------------------------------------------------------------------


def _processor_image_patch_size(processor: Any) -> int:
    image_processor = getattr(processor, "image_processor", None)
    patch_size = getattr(image_processor, "patch_size", None)
    if patch_size is None:
        return 14
    try:
        return int(patch_size)
    except (TypeError, ValueError):
        return 14


def _processor_image_pixel_kwargs(processor: Any) -> dict[str, int]:
    image_processor = getattr(processor, "image_processor", None)
    min_pixels = getattr(image_processor, "min_pixels", None)
    max_pixels = getattr(image_processor, "max_pixels", None)
    kwargs: dict[str, int] = {}
    if min_pixels is not None and int(min_pixels) > 0:
        kwargs["min_pixels"] = int(min_pixels)
    if max_pixels is not None and int(max_pixels) > 0:
        kwargs["max_pixels"] = int(max_pixels)
    return kwargs


# ---------------------------------------------------------------------------
# Model + processor loading at module scope
# ---------------------------------------------------------------------------

from transformers import AutoProcessor, AutoModelForImageTextToText
import peft

print("Loading processor...")
processor = AutoProcessor.from_pretrained(
    BASE_MODEL,
    trust_remote_code=True,
    max_pixels=IMAGE_MAX_PIXELS,
)

# Chat-template fix: copy from tokenizer if processor lacks one
tokenizer = getattr(processor, "tokenizer", processor)
if getattr(processor, "chat_template", None) is None and getattr(
    tokenizer, "chat_template", None
):
    processor.chat_template = tokenizer.chat_template
if getattr(tokenizer, "pad_token_id", None) is None:
    tokenizer.pad_token = tokenizer.eos_token
if hasattr(tokenizer, "padding_side"):
    tokenizer.padding_side = "left"

print("Loading base model...")
model = AutoModelForImageTextToText.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    low_cpu_mem_usage=True,
)

print("Loading LoRA adapter (manual CPU load for ZeroGPU compatibility)...")
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import peft

adapter_config = peft.LoraConfig.from_pretrained(ADAPTER)
adapter_weights_path = hf_hub_download(ADAPTER, "adapter_model.safetensors")
adapter_state_dict = load_file(adapter_weights_path, device="cpu")
model = peft.get_peft_model(model, adapter_config)
peft.set_peft_model_state_dict(model, adapter_state_dict)
model = model.merge_and_unload()
model.eval()
model.to("cuda")
print("Model ready.")


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------


@spaces.GPU(duration=120)
def predict(
    image: str,
    max_new_tokens: int = MAX_NEW_TOKENS_DEFAULT,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> tuple[str, str, str]:
    """Extract Markush CXSMILES from a chemical structure image.

    Args:
        image: Filepath to the input chemical structure image.
        max_new_tokens: Maximum number of tokens to generate.
        progress: Gradio progress bar.

    Returns:
        Tuple of (raw_output, cxsmiles, stable_groups).
    """
    from qwen_vl_utils import process_vision_info

    image_patch_size = _processor_image_patch_size(processor)
    image_pixel_kwargs = _processor_image_pixel_kwargs(processor)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": str(image), **image_pixel_kwargs},
                {"type": "text", "text": PROMPT},
            ],
        }
    ]

    # Apply chat template
    try:
        text_prompt = processor.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,
        )
    except TypeError:
        try:
            text_prompt = processor.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )
        except TypeError:
            text_prompt = processor.apply_chat_template(messages, tokenize=False)

    imgs, _ = process_vision_info(
        messages,
        image_patch_size=image_patch_size,
    )

    inputs = processor(
        text=[text_prompt],
        images=imgs,
        return_tensors="pt",
        padding=True,
    )
    inputs = {
        k: v.to("cuda") if torch.is_tensor(v) else v for k, v in inputs.items()
    }
    prompt_len = inputs["input_ids"].shape[1]

    pad_token_id = getattr(tokenizer, "pad_token_id", None) or getattr(
        tokenizer, "eos_token_id", None
    )

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            do_sample=False,
            max_new_tokens=max_new_tokens,
            pad_token_id=pad_token_id,
            eos_token_id=[ENDOFTEXT_TOKEN_ID, IM_END_TOKEN_ID],
            use_cache=True,
        )

    generated = output_ids[0, prompt_len:]
    raw_text = tokenizer.decode(generated, skip_special_tokens=True)
    cleaned = clean_prediction(raw_text)

    # Extract CXSMILES and stable groups
    m = RE_MARKUSH.search(cleaned)
    if m:
        cxsmiles = m.group(1).strip()
        stable = (m.group(2) or "").strip()
    else:
        cxsmi_match = RE_CXSMI.search(cleaned)
        cxsmiles = cxsmi_match.group(1).strip() if cxsmi_match else cleaned
        stable_match = RE_STABLE.search(cleaned)
        stable = stable_match.group(1).strip() if stable_match else ""

    return cleaned, cxsmiles, stable


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # MarkushGlyph — Markush Structure Recognition

        **MarkushGlyph** is a vision-language model (LoRA fine-tune of Qwen3.5-2B-Base)
        that converts images of patent Markush chemical structures into
        [CXSMILES](https://docs.chemaxon.com/display/docs/chemaxon-ex
        -extended-smiles-and-smarts.md) strings.

        Upload a chemical structure image and the model will extract the Markush
        structure, returning the full XML output, the CXSMILES core, and the
        stable group definitions.

        [Model card](https://huggingface.co/EdisonScientific/MarkushGlyph) |
        [Code repository](https://github.com/EdisonScientific/glyph)
        """
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Chemical structure image",
                    type="filepath",
                    height=400,
                )
                run_btn = gr.Button("Extract Markush Structure", variant="primary")

                with gr.Accordion("Advanced settings", open=False):
                    max_tokens = gr.Slider(
                        label="Max new tokens",
                        minimum=128,
                        maximum=2048,
                        value=MAX_NEW_TOKENS_DEFAULT,
                        step=128,
                        info="Maximum number of tokens to generate for the output.",
                    )

            with gr.Column(scale=1):
                raw_output = gr.Textbox(
                    label="Raw model output (XML)",
                    lines=8,
                )
                cxsmiles_output = gr.Textbox(
                    label="CXSMILES",
                    lines=3,
                )
                stable_output = gr.Textbox(
                    label="Stable groups",
                    lines=4,
                )

    run_btn.click(
        fn=predict,
        inputs=[image_input, max_tokens],
        outputs=[raw_output, cxsmiles_output, stable_output],
        api_name="predict",
    )

    gr.Examples(
        examples=[
            ["markush_m2s_13.png"],
            ["imatinib.png"],
            ["cholesterol_ester.png"],
        ],
        inputs=[image_input],
        outputs=[raw_output, cxsmiles_output, stable_output],
        fn=predict,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)