"""TissueDecision-9B with llm2jev: one forward pass per question, label-token probabilities (no text generation).

Start the model first (any GPU with >= 24 GB):
  vllm serve lighteternal/tissuedecision-9b --max-logprobs 256 --return-tokens-as-token-ids --enable-scale-out \
      --enable-prefix-caching --max-model-len 16384 --default-chat-template-kwargs '{"enable_thinking": false}'
Then:  pip install "llm2jev[vision]==0.6.1"  and  python usage_example.py lesion_tile.jpg "Rat liver, H&E"
(`prepare_image` comes from tissuedecision_server.py in the model repo: it resizes to the training pixel budget.)
"""
import json
import os
import sys

from llm2jev import LLM2Jev
from llm2jev.backends import VLLM
from transformers import AutoProcessor

from tissuedecision_server import prepare_image

MODEL = os.environ.get("TISSUEDECISION_MODEL", "lighteternal/tissuedecision-9b")
QUESTIONS = {
    "lesion": {"type": "noul", "instructions": "Is there any histopathological lesion (abnormal finding) in this image?"},
    "finding": {"type": "choice", "instructions": "What is the primary histopathological finding in this image?",
                "criteria": {"necrosis": None, "single cell necrosis": None, "inflammatory cell infiltrate": None,
                             "hemorrhage": None, "vacuolation": None, "no significant findings": None}},
    "organ": {"type": "choice", "instructions": "Which organ is shown?",
              "criteria": {"liver": None, "kidney": None, "thyroid gland": None, "testis": None, "urinary bladder": None,
                           "lung": None, "heart": None, "stomach": None, "salivary gland": None}},
}

if __name__ == "__main__":
    image = sys.argv[1] if len(sys.argv) > 1 else "lesion_tile.jpg"
    context = sys.argv[2] if len(sys.argv) > 2 else None  # e.g. "Rat liver, H&E"; optional
    jev = LLM2Jev(AutoProcessor.from_pretrained(MODEL), VLLM("http://localhost:8000", MODEL), temperature=1.0)
    content = [{"type": "image", "image": prepare_image(image)}] + ([{"type": "text", "text": context}] if context else [])
    answers = jev([{"role": "user", "content": content}], QUESTIONS)  # raw probabilities (calibration: see the server)
    print(json.dumps(answers, indent=1))