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 ......; " "the 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".*?\s*", _XML_FLAGS) RE_CXSMI = re.compile(r"(.*?)", _XML_FLAGS) RE_STABLE = re.compile(r"(.*?)", _XML_FLAGS) RE_MARKUSH = re.compile( r"\s*(.*?)\s*(?:(.*?))?\s*", _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)