Update app.py

image → OCR (Chandra/Dots via radio) → take final Raw stream text → ClinicalNER → spell-check (TF-IDF / SymSpell / RapidFuzz via radio) → show final Markdown with scores.

It also dynamically pulls ner.py, tfidf_phonetic.py, symspell_matcher.py, rapidfuzz_matcher.py, and your drug_dictionary.csv from your private HF repo

app.py

@@ -1,7 +1,385 @@
+import os
+import time
+from threading import Thread
+from typing import Iterable, Dict, Any, Optional, List
+
 import gradio as gr
+import spaces
+import torch
+from PIL import Image
+
+from transformers import (
+    Qwen3VLForConditionalGeneration,
+    AutoModelForCausalLM,
+    AutoProcessor,
+    TextIteratorStreamer,
+)
+
+from gradio.themes import Soft
+from gradio.themes.utils import colors, fonts, sizes
+
+# -----------------------------
+# Private repo: dynamic import
+# -----------------------------
+import importlib.util
+from huggingface_hub import hf_hub_download
+
+REPO_ID = "IFMedTech/new_model_private"  # your private backend repo
+
+# Map filenames to exported class names
+PY_MODULES = {
+    "ner.py": "ClinicalNER",
+    "tfidf_phonetic.py": "TfidfPhoneticMatcher",
+    "symspell_matcher.py": "SymSpellMatcher",
+    "rapidfuzz_matcher.py": "RapidFuzzMatcher",
+    # 'drug_dictionary.csv' is data, not a module
+}
+
+HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+
+def _dynamic_import(module_path: str, class_name: str):
+    spec = importlib.util.spec_from_file_location(class_name, module_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)  # type: ignore
+    return getattr(module, class_name)
+
+# Load all private classes and data (CSV)
+priv_classes: Dict[str, Any] = {}
+drug_csv_path: Optional[str] = None
+try:
+    if HF_TOKEN is None:
+        print("[Private] WARNING: HUGGINGFACE_TOKEN not set; NER/Spell-check will be unavailable.")
+    else:
+        for fname, cls in PY_MODULES.items():
+            path = hf_hub_download(repo_id=REPO_ID, filename=fname, token=HF_TOKEN)
+            if cls:
+                priv_classes[cls] = _dynamic_import(path, cls)
+                print(f"[Private] Loaded class: {cls} from {fname}")
+        # Download the drug dictionary CSV
+        drug_csv_path = hf_hub_download(repo_id=REPO_ID, filename="drug_dictionary.csv", token=HF_TOKEN)
+        print(f"[Private] Downloaded CSV at: {drug_csv_path}")
+except Exception as e:
+    print(f"[Private] ERROR loading private backend: {e}")
+    priv_classes = {}
+    drug_csv_path = None
+
+# ----------------------------
+# THEME
+# ----------------------------
+colors.steel_blue = colors.Color(
+    name="steel_blue",
+    c50="#EBF3F8",
+    c100="#D3E5F0",
+    c200="#A8CCE1",
+    c300="#7DB3D2",
+    c400="#529AC3",
+    c500="#4682B4",
+    c600="#3E72A0",
+    c700="#36638C",
+    c800="#2E5378",
+    c900="#264364",
+    c950="#1E3450",
+)
+
+class SteelBlueTheme(Soft):
+    def __init__(
+        self,
+        *,
+        primary_hue: colors.Color | str = colors.gray,
+        secondary_hue: colors.Color | str = colors.steel_blue,
+        neutral_hue: colors.Color | str = colors.slate,
+        text_size: sizes.Size | str = sizes.text_lg,
+        font: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
+        ),
+        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
+        ),
+    ):
+        super().__init__(
+            primary_hue=primary_hue,
+            secondary_hue=secondary_hue,
+            neutral_hue=neutral_hue,
+            text_size=text_size,
+            font=font,
+            font_mono=font_mono,
+        )
+        super().set(
+            background_fill_primary="*primary_50",
+            background_fill_primary_dark="*primary_900",
+            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
+            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
+            button_primary_text_color="white",
+            button_primary_text_color_hover="white",
+            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
+            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_800)",
+            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_500)",
+            button_secondary_text_color="black",
+            button_secondary_text_color_hover="white",
+            button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
+            button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
+            button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
+            button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
+            slider_color="*secondary_500",
+            slider_color_dark="*secondary_600",
+            block_title_text_weight="600",
+            block_border_width="3px",
+            block_shadow="*shadow_drop_lg",
+            button_primary_shadow="*shadow_drop_lg",
+            button_large_padding="11px",
+            color_accent_soft="*primary_100",
+            block_label_background_fill="*primary_200",
+        )
+
+steel_blue_theme = SteelBlueTheme()
+
+css = """
+#main-title h1 { font-size: 2.3em !important; }
+#output-title h2 { font-size: 2.1em !important; }
+"""
+
+# ----------------------------
+# RUNTIME / DEVICE
+# ----------------------------
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
+print("CUDA_VISIBLE_DEVICES =", os.environ.get("CUDA_VISIBLE_DEVICES"))
+print("torch.__version__     =", torch.__version__)
+print("torch.version.cuda    =", torch.version.cuda)
+print("cuda available        =", torch.cuda.is_available())
+print("cuda device count     =", torch.cuda.device_count())
+if torch.cuda.is_available():
+    print("using device          =", torch.cuda.get_device_name(0))
+
+use_cuda = torch.cuda.is_available()
+device = torch.device("cuda:0" if use_cuda else "cpu")
+if use_cuda:
+    torch.backends.cudnn.benchmark = True
+
+DTYPE_FP16 = torch.float16 if use_cuda else torch.float32
+DTYPE_BF16 = torch.bfloat16 if use_cuda else torch.float32
+
+# ----------------------------
+# MODELS: Chandra-OCR + Dots.OCR
+# ----------------------------
+# 1) Chandra-OCR (Qwen3VL)
+MODEL_ID_V = "datalab-to/chandra"
+processor_v = AutoProcessor.from_pretrained(MODEL_ID_V, trust_remote_code=True)
+model_v = Qwen3VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_V, trust_remote_code=True, torch_dtype=DTYPE_FP16
+).to(device).eval()
+
+# 2) Dots.OCR (flash_attn2 if available, else SDPA)
+MODEL_PATH_D = "prithivMLmods/Dots.OCR-Latest-BF16"
+processor_d = AutoProcessor.from_pretrained(MODEL_PATH_D, trust_remote_code=True)
+attn_impl = "sdpa"
+try:
+    import flash_attn  # noqa: F401
+    if use_cuda:
+        attn_impl = "flash_attention_2"
+except Exception:
+    attn_impl = "sdpa"
+
+model_d = AutoModelForCausalLM.from_pretrained(
+    MODEL_PATH_D,
+    attn_implementation=attn_impl,
+    torch_dtype=DTYPE_BF16,
+    device_map="auto" if use_cuda else None,
+    trust_remote_code=True
+).eval()
+if not use_cuda:
+    model_d.to(device)
+
+# ----------------------------
+# GENERATION (OCR → NER → Spell-check)
+# ----------------------------
+MAX_MAX_NEW_TOKENS = 4096
+DEFAULT_MAX_NEW_TOKENS = 2048
+
+@spaces.GPU
+def generate_image(model_name: str,
+                   text: str,
+                   image: Image.Image,
+                   max_new_tokens: int,
+                   temperature: float,
+                   top_p: float,
+                   top_k: int,
+                   repetition_penalty: float,
+                   spell_algo: str):
+    """
+    1) Streams OCR tokens to Raw output (unchanged).
+    2) After stream ends, runs ClinicalNER on FINAL RAW text → list[str] meds.
+    3) Runs selected spell-check approach on that list (top-5 with scores).
+    4) Markdown shows: OCR text + Clinical NER + Spell-check suggestions.
+    """
+    if image is None:
+        yield "Please upload an image.", "Please upload an image."
+        return
+
+    if model_name == "Chandra-OCR":
+        processor, model = processor_v, model_v
+    elif model_name == "Dots.OCR":
+        processor, model = processor_d, model_d
+    else:
+        yield "Invalid model selected.", "Invalid model selected."
+        return
+
+    # Prepare prompt
+    messages = [{
+        "role": "user",
+        "content": [
+            {"type": "image"},
+            {"type": "text", "text": text},
+        ]
+    }]
+    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    # Preprocess
+    inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True)
+    inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}
+
+    # Streamer
+    tokenizer = getattr(processor, "tokenizer", None) or processor
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    gen_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+    )
+
+    # Start generation
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    # 1) live OCR stream to Raw & Markdown mirrors during stream
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text.replace("<|im_end|>", "")
+        time.sleep(0.01)
+        yield buffer, buffer
+
+    # >>> Final RAW text for downstream steps <<<
+    final_ocr_text = buffer
+
+    # 2) ClinicalNER
+    try:
+        if "ClinicalNER" in priv_classes:
+            ClinicalNER = priv_classes["ClinicalNER"]
+            ner = ClinicalNER(token=HF_TOKEN)  # you can pass model_id=... if custom
+            meds: List[str] = ner(final_ocr_text) or []
+        else:
+            meds = []
+            print("[NER] ClinicalNER not available.")
+    except Exception as e:
+        meds = []
+        print(f"[NER] Error running ClinicalNER: {e}")
+
+    # Build Markdown with OCR + NER section
+    md = final_ocr_text
+    md += "\n\n---\n### Clinical NER (Medications)\n"
+    if meds:
+        for m in meds:
+            md += f"- {m}\n"
+    else:
+        md += "- None detected\n"
+
+    # 3) Spell-check on NER output using selected approach
+    spell_section = "\n---\n### Spell-check suggestions (" + spell_algo + ")\n"
+    corr: Dict[str, List] = {}
+    try:
+        if meds and drug_csv_path:
+            if spell_algo == "TF-IDF + Phonetic" and "TfidfPhoneticMatcher" in priv_classes:
+                Cls = priv_classes["TfidfPhoneticMatcher"]
+                checker = Cls(csv_path=drug_csv_path, column="drug_name", ngram_size=3, phonetic_weight=0.4)
+                corr = checker.match_list(meds, top_k=5, tfidf_threshold=0.15)
+
+            elif spell_algo == "SymSpell" and "SymSpellMatcher" in priv_classes:
+                Cls = priv_classes["SymSpellMatcher"]
+                checker = Cls(csv_path=drug_csv_path, column="drug_name", max_edit=2, prefix_len=7)
+                corr = checker.match_list(meds, top_k=5, min_score=0.4)
+
+            elif spell_algo == "RapidFuzz" and "RapidFuzzMatcher" in priv_classes:
+                Cls = priv_classes["RapidFuzzMatcher"]
+                checker = Cls(csv_path=drug_csv_path, column="drug_name")
+                corr = checker.match_list(meds, top_k=5, threshold=70.0)
+            else:
+                spell_section += "- Spell-check backend unavailable.\n"
+        else:
+            spell_section += "- No NER output or dictionary missing.\n"
+    except Exception as e:
+        spell_section += f"- Spell-check error: {e}\n"
+
+    # Format suggestions
+    if corr:
+        for raw in meds:
+            cand = corr.get(raw, [])
+            if cand:
+                spell_section += f"- **{raw}**\n"
+                for c, s in cand:
+                    spell_section += f"  - {c} (score={s:.3f})\n"
+            else:
+                spell_section += f"- **{raw}**\n  - (no suggestions)\n"
+
+    final_md = md + spell_section
+
+    # 4) Final yield: raw unchanged; Markdown has NER + spell-check
+    yield final_ocr_text, final_md
+
+# ----------------------------
+# UI
+# ----------------------------
+image_examples = [
+    ["OCR the content perfectly.", "examples/3.jpg"],
+    ["Perform OCR on the image.", "examples/1.jpg"],
+    ["Extract the contents. [page].", "examples/2.jpg"],
+]
+
+with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
+    gr.Markdown("# **Multimodal OCR3**", elem_id="main-title")
+    with gr.Row():
+        with gr.Column(scale=2):
+            image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
+            image_upload = gr.Image(type="pil", label="Upload Image", height=290)
+            image_submit = gr.Button("Submit", variant="primary")
+            gr.Examples(examples=image_examples, inputs=[image_query, image_upload])
+
+            # Spell-check selection
+            spell_choice = gr.Radio(
+                choices=["TF-IDF + Phonetic", "SymSpell", "RapidFuzz"],
+                label="Select Spell-check Approach",
+                value="TF-IDF + Phonetic"
+            )
+
+            with gr.Accordion("Advanced options", open=False):
+                max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
+                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.7)
+                top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
+                top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
+                repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.1)
+
+        with gr.Column(scale=3):
+            gr.Markdown("## Output", elem_id="output-title")
+            output = gr.Textbox(label="Raw Output Stream", interactive=False, lines=11, show_copy_button=True)
+            with gr.Accordion("(Result.md)", open=False):
+                markdown_output = gr.Markdown(label="(Result.Md)")
+
+            model_choice = gr.Radio(
+                choices=["Chandra-OCR", "Dots.OCR"],
+                label="Select OCR Model",
+                value="Chandra-OCR"
+            )
 
-def greet(name):
-    return "Hello " + name + "!!"
+    image_submit.click(
+        fn=generate_image,
+        inputs=[model_choice, image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty, spell_choice],
+        outputs=[output, markdown_output]
+    )
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+if __name__ == "__main__":
+    demo.queue(max_size=50).launch(mcp_server=True, ssr_mode=False, show_error=True)