Blame for janus_engine.py

import gc
import logging
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Dict, List

import torch

def _log_runtime_versions():
    try:
        import transformers
        logger.info("Transformers version: %s", transformers.__version__)
    except Exception as exc:
        logger.warning("Could not read Transformers version: %s", exc)
    logger.info("Torch version: %s | CUDA available: %s", torch.__version__, torch.cuda.is_available())


logger = logging.getLogger("featurex-janus")

MODEL_ID = os.environ.get("JANUS_MODEL_ID", "ZrH42/Janus-Pro-CXR-Final")
REPO_URL = "https://github.com/ZrH42/Janus-Pro-CXR.git"
REPO_DIR = Path(os.environ.get("JANUS_REPO_DIR", "/tmp/Janus-Pro-CXR"))
JANUS_ENABLED = os.environ.get("JANUS_ENABLED", "true").strip().lower() in {"1", "true", "yes", "on"}
BILINGUAL = os.environ.get("JANUS_BILINGUAL", "false").strip().lower() in {"1", "true", "yes", "on"}

_PROCESSOR = None
_MODEL = None
_INITIALIZED = False


def _clone_repo():
    if (REPO_DIR / "janus").exists():
        return
    REPO_DIR.parent.mkdir(parents=True, exist_ok=True)
    logger.info("Cloning Janus-Pro-CXR repository to %s", REPO_DIR)
    subprocess.run(
        ["git", "clone", "--depth", "1", REPO_URL, str(REPO_DIR)],
        check=True,
        stdout=subprocess.DEVNULL,
    )


def _patch_compatibility():
    siglip = REPO_DIR / "janus" / "models" / "siglip_vit.py"
    if siglip.exists():
        text = siglip.read_text(encoding="utf-8")
        old = "dpr = [drop_path_rate * i / (depth - 1) for i in torch.linspace(0, depth - 1, depth)]"
        if old in text:
            text = text.replace(old, "dpr = [drop_path_rate * i / max(depth - 1, 1) for i in range(depth)]")
            siglip.write_text(text, encoding="utf-8")

    if str(REPO_DIR) not in sys.path:
        sys.path.insert(0, str(REPO_DIR))

    try:
        from janus.models.modeling_vlm import MultiModalityCausalLM
        if not hasattr(MultiModalityCausalLM, "all_tied_weights_keys"):
            MultiModalityCausalLM.all_tied_weights_keys = property(lambda self: {})
    except Exception:
        pass


def _initialize():
    global _PROCESSOR, _MODEL, _INITIALIZED
    if _INITIALIZED:
        return
    _log_runtime_versions()
    if not JANUS_ENABLED:
        _INITIALIZED = True
        return
    if not torch.cuda.is_available():
        raise RuntimeError("Janus requires CUDA/ZeroGPU. The analyze function must execute on a ZeroGPU Space.")

    _clone_repo()
    _patch_compatibility()

    from transformers import AutoModelForCausalLM
    from janus.models import VLChatProcessor

    logger.info("Loading Janus processor: %s", MODEL_ID)
    _PROCESSOR = VLChatProcessor.from_pretrained(MODEL_ID)
    logger.info("Loading Janus model: %s", MODEL_ID)
    _MODEL = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        low_cpu_mem_usage=False,
    ).to(torch.bfloat16).cuda().eval()
    _INITIALIZED = True
    logger.info("Janus model ready on CUDA")


def _prompt(language: str) -> str:
    if language == "ar":
        return (
            "حلل صورة أشعة الصدر هذه كمساعد في طب الأشعة. اكتب تقريراً طبياً موجزاً باللغة العربية فقط. "
            "استخدم العناوين العربية التالية عند الحاجة: النتائج، الانطباع، التوصيات. "
            "يجب أن يكون نص التقرير نفسه بالعربية بالكامل، ولا تكتب FINDINGS أو IMPRESSION أو RECOMMENDATIONS، "
            "ولا تعِد كتابة التقرير بالإنجليزية. صف فقط العلامات الشعاعية المرئية. "
            "لا تخترع أجهزة أو قياسات أو فحوصاً سابقة أو علاجاً أو تاريخاً مرضياً. "
            "استخدم صياغة حذرة عند وجود عدم يقين. احتفظ بالأرقام والاختصارات الطبية الضرورية فقط."
        )
    return (
        "Analyze this chest X-ray as a radiology assistant. Return a concise report using these headings "
        "exactly when applicable: FINDINGS, IMPRESSION, RECOMMENDATIONS. "
        "Describe only visible radiographic findings. Do not invent devices, measurements, prior studies, "
        "treatment, or patient history. Use cautious wording for uncertain findings. Write the report in English."
    )


def _extract_text(result) -> str:
    if isinstance(result, str):
        return result.strip()
    if isinstance(result, dict):
        for key in ("text", "answer", "response", "generated_text"):
            if result.get(key):
                return str(result[key]).strip()
    return str(result).strip()


def _generate_one(image, language: str) -> str:
    from janus.utils.io import load_pil_images

    # Janus-Pro-CXR's official inference helper accepts an image path.
    # Keep the request image isolated in a temporary file and delete it after generation.
    suffix = ".png"
    tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    tmp_path = tmp.name
    tmp.close()
    image.save(tmp_path, format="PNG")

    conversation = [
        {
            "role": "User",
            "content": f"<image_placeholder>\n{_prompt(language)}",
            "images": [tmp_path],
        },
        {"role": "Assistant", "content": ""},
    ]

    pil_images = load_pil_images(conversation)
    prepare = _PROCESSOR(
        conversations=conversation,
        images=pil_images,
        force_batchify=True,
    ).to(_MODEL.device)

    inputs_embeds = _MODEL.prepare_inputs_embeds(**prepare)
    outputs = _MODEL.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare.attention_mask,
        pad_token_id=_PROCESSOR.tokenizer.eos_token_id,
        bos_token_id=_PROCESSOR.tokenizer.bos_token_id,
        eos_token_id=_PROCESSOR.tokenizer.eos_token_id,
        max_new_tokens=512,
        do_sample=False,
        use_cache=True,
        output_scores=True,
        return_dict_in_generate=True,
    )
    try:
        return _extract_text(_PROCESSOR.tokenizer.decode(outputs.sequences[0].cpu().tolist(), skip_special_tokens=True))
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass


def _arabic_char_count(text: str) -> int:
    return len(re.findall(r"[\u0600-\u06FF]", text or ""))


def _latin_char_count(text: str) -> int:
    return len(re.findall(r"[A-Za-z]", text or ""))


def _has_arabic(text: str) -> bool:
    # Require actual Arabic content, not a single Arabic punctuation/character.
    return _arabic_char_count(text) >= 8


def _has_english_letters(text: str) -> bool:
    return _latin_char_count(text) >= 8


def _is_predominantly_arabic(text: str) -> bool:
    ar = _arabic_char_count(text)
    en = _latin_char_count(text)
    return ar >= 8 and ar >= max(8, en * 1.25)


def _is_predominantly_english(text: str) -> bool:
    ar = _arabic_char_count(text)
    en = _latin_char_count(text)
    return en >= 8 and en >= max(8, ar * 1.25)


def _generate_arabic_verified(image, attempts: int = 2) -> str:
    """Generate Arabic and return an empty string unless the result is genuinely Arabic.

    This is important: returning the last English candidate here used to leak English
    into ``report_ar_raw`` and then into the Arabic ResultScreen card.
    """
    for attempt in range(1, attempts + 1):
        candidate = _generate_one(image, "ar")
        if _is_predominantly_arabic(candidate):
            logger.info("Janus Arabic report verified on attempt %d", attempt)
            return candidate.strip()
        logger.warning(
            "Janus Arabic attempt %d failed language validation (ar=%d, en=%d); retrying",
            attempt, _arabic_char_count(candidate), _latin_char_count(candidate),
        )
    # Never expose an unverified candidate as Arabic. inference.py will use the
    # deterministic Arabic template instead.
    return ""


def generate_reports(image, requested_language: str = "ar") -> Dict[str, object]:
    """Generate independently generated, language-validated Arabic and English reports."""
    if not JANUS_ENABLED:
        return {"enabled": False, "report_ar": "", "report_en": "", "requested_raw_text": ""}

    _initialize()
    reports = {
        "enabled": True,
        "report_ar": "",
        "report_en": "",
        "requested_raw_text": "",
        "arabic_verified": False,
        "english_verified": False,
    }
    with torch.inference_mode():
        english_candidate = _generate_one(image, "en")
        if _is_predominantly_english(english_candidate):
            reports["report_en"] = english_candidate.strip()
            reports["english_verified"] = True
        else:
            logger.warning(
                "Janus English output failed language validation (ar=%d, en=%d); using empty raw text",
                _arabic_char_count(english_candidate), _latin_char_count(english_candidate),
            )

        if BILINGUAL:
            reports["report_ar"] = _generate_arabic_verified(image, attempts=2)
            reports["arabic_verified"] = _is_predominantly_arabic(reports["report_ar"])

    # requested_raw_text is always language-safe. It is never a fallback from one
    # language into the other.
    reports["requested_raw_text"] = (
        reports["report_ar"] if requested_language == "ar" and reports["arabic_verified"]
        else reports["report_en"] if requested_language == "en" and reports["english_verified"]
        else ""
    )
    return reports


def parse_radiology_text(text: str) -> Dict[str, object]:
    text = (text or "").strip()
    if not text:
        return {"findings": "", "impression": "", "recommendations": [], "raw_text": ""}

    # Normalize common heading variants without trying to reinterpret medical content.
    normalized = re.sub(r"\r\n?", "\n", text)
    pattern = re.compile(
        r"(?:^|[\n\r]|[.\u061b;])\s*"
        r"(FINDINGS?|IMPRESSION|CONCLUSION|RECOMMENDATIONS?|RECOMMENDATION|النتائج|الملاحظات|الانطباع|التوصيات)"
        r"\s*:?\s*",
        flags=re.IGNORECASE,
    )
    matches = list(pattern.finditer(normalized))
    sections: Dict[str, str] = {}
    for i, match in enumerate(matches):
        raw_heading = match.group(1).lower()
        start = match.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(normalized)
        body = normalized[start:end].strip(" \n:-•")
        if raw_heading.startswith("finding") or raw_heading in {"النتائج", "الملاحظات"}:
            sections["findings"] = body
        elif raw_heading in {"impression", "conclusion", "الانطباع"}:
            sections["impression"] = body
        else:
            sections["recommendations"] = body

    if not sections:
        sections["findings"] = normalized
        sections["impression"] = normalized

    recs = []
    if sections.get("recommendations"):
        for line in re.split(r"\n+", sections["recommendations"]):
            line = re.sub(r"^\s*(?:[-*•]|\d+[.)])\s*", "", line).strip()
            if line:
                recs.append(line)

    return {
        "findings": sections.get("findings", ""),
        "impression": sections.get("impression", ""),
        "recommendations": recs,
        "raw_text": text,
    }