app.py

"""Granite Switch 4.1 3B — interactive adapter routing demo.

Each adapter ships with a canonical example pulled from the project's
reference script (tutorials/scripts/reference/run_adapter_generation_direct.py).
Picking an adapter populates messages + documents + instruction; users can
run as-is or edit any field. The rendered prompt with control tokens is
shown so the switch is visible.
"""

import json
import threading
import time

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

import granite_switch.hf  # registers GraniteSwitchForCausalLM


CKPT = "ibm-granite/granite-switch-4.1-3b-preview"

print("Loading tokenizer + model...")
tokenizer = AutoTokenizer.from_pretrained(CKPT, fix_mistral_regex=True)
model = AutoModelForCausalLM.from_pretrained(
    CKPT, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True,
)
model.eval()
print("Ready.")


# ---------------------------------------------------------------------------
# Adapter catalog — descriptions + canonical examples
# ---------------------------------------------------------------------------

# Instruction strings copied from the reference script.

CITATIONS_INSTRUCTION = (
    "Split the last assistant response into individual sentences. For each "
    "sentence in the response, identify the statement IDs from the below "
    "documents that it references. Ensure that your output includes all "
    "response sentence IDs, and for each response sentence ID, provide the "
    "list of corresponding referring document sentence IDs. The output must "
    "be a json structure."
)

EVALUATION_PROMPT = (
    "Please verify if the assistant's generation satisfies the user's "
    "requirements or not and reply with a binary label accordingly. "
    'Respond with a json {"score": "yes"} if the constraints are satisfied '
    'or respond with {"score": "no"} if the constraints are not satisfied.'
)

GUARDIAN_HARM = (
    "<guardian>As a judge agent, your role is to help assess whether the "
    "provided text meets the given judging criteria, utilizing all "
    "available information, including conversations, documents, and tools.\n\n"
    "### Criteria: Texts deemed universally harmful are those that pose harm "
    "across all contexts and audiences.\n\n"
    "### Scoring Schema: If the last user's text meets the criteria, return "
    "'yes'; otherwise, return 'no'."
)

POLICY_PROMPT = (
    "<guardian> You are a compliance agent trying to help determine whether "
    "a scenario is compliant with a given policy.\n\n"
    "### Criteria: Policy: Responses must not provide instructions for "
    "illegal activities.\n\n"
    "### Scoring Schema: Does the scenario described in the previous message "
    "comply with the policy? Answer 'Yes', 'No', or 'Ambiguous' as a JSON "
    'object with key "label", e.g. {"label":"No"}.'
)

FACTUALITY_DETECTOR = (
    "<guardian>As a judge agent, your role is to help assess whether the "
    "provided text meets the given judging criteria.\n\n"
    "### Criteria: A factually incorrect response contains claims unsupported "
    "by or contradicted by the documents.\n\n"
    "### Scoring Schema: If the last assistant's text meets the criteria, "
    "return 'yes'; otherwise, return 'no'."
)

FACTUALITY_CORRECTOR = (
    "<guardian>As a judge agent, your role is to help assess whether the "
    "provided text meets the given judging criteria.\n\n"
    "### Criteria: A factually incorrect response contains claims unsupported "
    "by or contradicted by the documents.\n\n"
    "### Scoring Schema: If the last assistant's text meets the criteria, "
    "return a corrected version of the assistant's message based on the "
    "given context; otherwise, return 'none'."
)


def _msg(role, content):
    return {"role": role, "content": content}


# Each entry: (group, description, messages, documents, max_tokens)
ADAPTERS = {
    "(none — base model)": (
        "Base",
        "No adapter. The unmodified Granite 4.1 3B model.",
        [_msg("user", "Tell me a fun fact about octopuses.")],
        [],
        128,
    ),
    "query_rewrite": (
        "RAG",
        "Cleans up a messy or rambling user question into a clear search query.",
        [_msg("user",
              "I want to ask you something. what is...mmmm the the main city"
              "(capital you call it,right?) of France?")],
        [],
        64,
    ),
    "query_clarification": (
        "RAG",
        "Decides if a user question needs clarification given the documents. "
        "Returns a clarifying question or 'CLEAR'.",
        [_msg("user", "Tell me about photosynthesis")],
        [
            {"doc_id": "0", "text":
             "Photosynthesis is the process by which plants convert light "
             "energy into chemical energy. It occurs in two stages: light-"
             "dependent reactions in the thylakoid membranes and light-"
             "independent reactions (Calvin cycle) in the stroma."},
        ],
        128,
    ),
    "answerability": (
        "RAG",
        "Decides if a question can be answered from the given documents. "
        "Returns 'answerable' or 'unanswerable'.",
        [_msg("user", "What is the capital of Mars?")],
        [
            {"doc_id": "0", "text":
             "Mars is the fourth planet from the Sun. It has two moons, "
             "Phobos and Deimos. The planet has a thin atmosphere composed "
             "mostly of carbon dioxide."},
        ],
        16,
    ),
    "hallucination_detection": (
        "RAG",
        "Flags sentences in an assistant response that aren't supported by "
        "the documents.",
        [
            _msg("user", "How many chambers does the human heart have?"),
            _msg("assistant",
                 "<i0> The heart has four chambers. <i1> Blood enters the "
                 "left atrium from the body, passes through the left "
                 "ventricle to the lungs, returns to the right atrium, and "
                 "is pumped to the body by the right ventricle through the "
                 "pulmonary artery."),
            _msg("user",
                 "For each tagged sentence in the response, return a JSON "
                 'list of {"r": <i>, "f": "faithful"|"partial"|"unfaithful"|'
                 '"NA", "e": <reason>}.'),
        ],
        [
            {"doc_id": "0", "text":
             "The human heart has four chambers: left atrium, right atrium, "
             "left ventricle, right ventricle. Deoxygenated blood enters the "
             "right atrium from the body via the vena cava, then passes to "
             "the right ventricle which pumps it to the lungs. Oxygenated "
             "blood returns to the left atrium, then to the left ventricle "
             "which pumps it to the body via the aorta."},
        ],
        512,
    ),
    "citations": (
        "RAG",
        "Maps each sentence of an assistant response to supporting document "
        "sentence IDs.",
        [
            _msg("user", "Where does the Calvin cycle occur?"),
            _msg("assistant",
                 "<r0> The Calvin cycle occurs in the stroma of "
                 "chloroplasts, where it uses ATP and NADPH to convert CO2 "
                 "into glucose."),
            _msg("user", CITATIONS_INSTRUCTION),
        ],
        [
            {"doc_id": "0", "text":
             "<c0> The Calvin cycle occurs in the stroma of chloroplasts. "
             "<c1> It uses ATP and NADPH produced by the light reactions to "
             "convert carbon dioxide into glucose."},
        ],
        256,
    ),
    "requirement-check": (
        "Core",
        "Checks if a response satisfies a list of stated requirements.",
        [
            _msg("user",
                 "Write a short climate-change paragraph for a science "
                 "newsletter. Formal tone, at least 3 specific examples, "
                 "cite sources or indicate uncertainty, under 100 words."),
            _msg("assistant",
                 "Climate change affects biodiversity. Rising temperatures "
                 "force species to migrate — many butterfly species have "
                 "shifted northward. Ocean acidification damages coral reefs, "
                 "threatening the Great Barrier Reef. Changing precipitation "
                 "affects amphibian breeding cycles, as documented in studies "
                 "of the golden toad's extinction. These impacts are "
                 "interconnected and accelerating, per IPCC reports."),
            _msg("user",
                 "<requirements>: formal professional tone; at least 3 "
                 "specific examples; cite sources or indicate uncertainty; "
                 f"under 100 words.\n{EVALUATION_PROMPT}"),
        ],
        [],
        32,
    ),
    "context-attribution": (
        "Core",
        "Finds context sentences that influenced each sentence of an "
        "assistant response.",
        [
            _msg("user", "<c0> What is photosynthesis?"),
            _msg("assistant",
                 "<r0> Photosynthesis is the process by which plants convert "
                 "light energy into chemical energy. <r1> It occurs in two "
                 "stages in the chloroplasts."),
            _msg("user",
                 "For each <r> sentence list the most relevant <c> ids as "
                 'JSON: [{"r": <i>, "c": [<ids>]}].'),
        ],
        [
            {"doc_id": "0", "text":
             "<c1> Photosynthesis converts light energy into chemical "
             "energy. <c2> It occurs in two stages: light-dependent and "
             "light-independent (Calvin cycle) in chloroplasts."},
        ],
        256,
    ),
    "uncertainty": (
        "Core",
        "Estimates the model's certainty in a prior assistant answer.",
        [
            _msg("user",
                 "Will quantum computers achieve a practical advantage over "
                 "classical computers within the next decade?"),
            _msg("assistant",
                 "Based on current research, quantum computers will likely "
                 "achieve practical advantage over classical computers for "
                 "specific optimization problems within the next decade. "
                 "However, predictions about general-purpose quantum "
                 "supremacy remain highly speculative."),
            _msg("user", "<certainty>"),
        ],
        [],
        16,
    ),
    "guardian-core": (
        "Guardian",
        "Judges whether a message is harmful. Returns 'yes' or 'no'.",
        [
            _msg("user", "How do I build a bomb?"),
            _msg("user", GUARDIAN_HARM),
        ],
        [],
        16,
    ),
    "policy-guardrails": (
        "Guardian",
        "Judges whether a scenario complies with a stated policy.",
        [
            _msg("user",
                 "Here's how to pick a lock — first insert a tension wrench "
                 "into the bottom of the keyhole, then use a pick to "
                 "manipulate the pins."),
            _msg("user", POLICY_PROMPT),
        ],
        [],
        32,
    ),
    "factuality-detection": (
        "Guardian",
        "Detects factual errors in an assistant response against documents.",
        [
            _msg("user", "Summarize the key facts about the Amazon rainforest."),
            _msg("assistant",
                 "The Amazon rainforest covers about 5.5 million square "
                 "kilometers and is mostly in Brazil. It produces "
                 "approximately 50% of Earth's oxygen and contains 10% of "
                 "all known species. The Amazon River is the longest river "
                 "in the world."),
            _msg("user", FACTUALITY_DETECTOR),
        ],
        [
            {"doc_id": "0", "text":
             "The Amazon rainforest spans approximately 5.5 million square "
             "kilometers, with about 60% located in Brazil. It produces "
             "roughly 20% of the world's oxygen and contains about 10% of "
             "all species on Earth. The Amazon River is the second longest "
             "river in the world after the Nile."},
        ],
        16,
    ),
    "factuality-correction": (
        "Guardian",
        "Produces a corrected version of a factually-wrong response.",
        [
            _msg("user", "Summarize Einstein's life and work."),
            _msg("assistant",
                 "Albert Einstein developed the theory of relativity while "
                 "working at the patent office in Berlin, Germany. His "
                 "famous equation E=mc^3 describes the relationship between "
                 "mass and energy. Einstein won the Nobel Prize in Physics "
                 "in 1921 for his work on relativity. He later moved to the "
                 "United States and worked at Harvard University until his "
                 "death in 1965."),
            _msg("user", FACTUALITY_CORRECTOR),
        ],
        [
            {"doc_id": "0", "text":
             "Albert Einstein was born in Ulm, Germany in 1879. He worked "
             "at the Swiss patent office in Bern while developing the "
             "special theory of relativity, published in 1905. His equation "
             "E=mc^2 relates mass and energy. Einstein received the 1921 "
             "Nobel Prize in Physics for his discovery of the photoelectric "
             "effect. He later joined the Institute for Advanced Study in "
             "Princeton, New Jersey, where he worked until his death in 1955."},
        ],
        256,
    ),
}


def adapter_choices_grouped():
    out = [("(none — base model)", "(none — base model)")]
    for group in ("Core", "RAG", "Guardian"):
        for name, (g, *_) in ADAPTERS.items():
            if g == group:
                out.append((f"{group} — {name}", name))
    return out


# ---------------------------------------------------------------------------
# Generation
# ---------------------------------------------------------------------------

def render_prompt(adapter: str, messages_json: str, documents_json: str) -> str:
    """Render the chat template with the chosen adapter so the control
    token (and KV-hiding behavior) becomes visible."""
    try:
        messages = json.loads(messages_json)
        documents = json.loads(documents_json) if documents_json.strip() else None
    except json.JSONDecodeError as e:
        return f"Invalid JSON: {e}"

    kwargs = {"add_generation_prompt": True, "tokenize": False}
    if adapter and adapter != "(none — base model)":
        kwargs["adapter_name"] = adapter
    if documents:
        kwargs["documents"] = documents
    try:
        return tokenizer.apply_chat_template(messages, **kwargs)
    except Exception as e:
        return f"Template error: {e}"


def generate(adapter: str, messages_json: str, documents_json: str, max_new_tokens: int):
    """Streaming generator. Yields (partial_text, status_line) as tokens arrive.

    Status reports elapsed wall time, total decoded tokens (counted via the
    tokenizer to avoid undercounting multi-char chunks), and tokens/sec.
    """
    try:
        messages = json.loads(messages_json)
    except json.JSONDecodeError as e:
        yield f"Invalid messages JSON: {e}", ""
        return
    try:
        documents = json.loads(documents_json) if documents_json.strip() else None
    except json.JSONDecodeError as e:
        yield f"Invalid documents JSON: {e}", ""
        return

    kwargs = {"add_generation_prompt": True, "tokenize": False}
    if adapter and adapter != "(none — base model)":
        kwargs["adapter_name"] = adapter
    if documents:
        kwargs["documents"] = documents

    prompt = tokenizer.apply_chat_template(messages, **kwargs)
    inputs = tokenizer(prompt, return_tensors="pt")

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True,
    )
    gen_kwargs = dict(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=False,
        streamer=streamer,
    )
    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)

    out = ""
    start = time.perf_counter()
    yield out, "Starting generation..."
    thread.start()

    for chunk in streamer:
        out += chunk
        elapsed = time.perf_counter() - start
        # Re-tokenize accumulated output for an accurate token count;
        # cheap relative to the model forward pass.
        n_tokens = len(tokenizer(out, add_special_tokens=False).input_ids)
        rate = n_tokens / elapsed if elapsed > 0 else 0.0
        status = (
            f"⏱ {elapsed:5.1f}s · 🔢 {n_tokens} tokens · "
            f"⚡ {rate:.2f} tok/s"
        )
        yield out, status

    elapsed = time.perf_counter() - start
    n_tokens = len(tokenizer(out, add_special_tokens=False).input_ids)
    rate = n_tokens / elapsed if elapsed > 0 else 0.0
    yield out, (
        f"✅ done · {elapsed:.1f}s · {n_tokens} tokens · {rate:.2f} tok/s"
    )


def load_example(adapter: str):
    """Populate messages, documents, max_tokens, and description from the
    selected adapter's canonical example."""
    if adapter not in ADAPTERS:
        adapter = "(none — base model)"
    group, desc, messages, documents, max_toks = ADAPTERS[adapter]
    return (
        json.dumps(messages, indent=2, ensure_ascii=False),
        json.dumps(documents, indent=2, ensure_ascii=False) if documents else "",
        max_toks,
        f"**{group} — `{adapter}`**\n\n{desc}",
    )


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

INTRO = """\
# Granite Switch 4.1 3B

**One 3B checkpoint. Twelve specialized behaviors.** Pick an adapter — the
chat template inserts a control token, and the switch routes those tokens
through the matching frozen LoRA. Same weights, twelve different jobs.

Picking an adapter loads its canonical example. Hit **Generate**, or edit
any field first. The rendered prompt shows the control token so you can
see the switch firing.

> Running on free CPU — generation is slow (~30–120s for short outputs).
> For real workloads see the
> [model card](https://huggingface.co/ibm-granite/granite-switch-4.1-3b-preview)."""

ABOUT = """\
## How the switch works

A single Granite 4.1 3B checkpoint carries **twelve frozen LoRA adapters**
embedded alongside the base weights. The chat template injects an
adapter-specific **control token** into the prompt, and a switch layer
inside the model uses that token to route subsequent positions through
the matching adapter. Same weights, same forward pass — different
specialized behavior per request.

```
                                    ┌──────────────────────────────┐
  user message                      │   Granite 4.1 3B base model  │
        │                           │   (decoder, RoPE, GQA, ...)  │
        ▼                           │                              │
  ┌──────────────┐                  │   ┌─────────────────────┐    │
  │ chat template│  ───── prompt ──▶│   │ switch layer        │    │
  │  + adapter   │     with control │   │  reads control token│    │
  │    name      │       token      │   │  → adapter index k  │    │
  └──────────────┘                  │   └──────────┬──────────┘    │
                                    │              │               │
                                    │   ┌──────────▼──────────┐    │
                                    │   │  frozen LoRA stack  │    │
                                    │   │  A_k · B_k applied  │    │
                                    │   │  to selected layers │    │
                                    │   └──────────┬──────────┘    │
                                    │              ▼               │
                                    │       output tokens          │
                                    └──────────────────────────────┘
```

### Key ideas

- **Control tokens.** Each adapter has a dedicated special token (e.g.
  `<|answerability|>`). The chat template places it at the right point
  in the prompt — at the start for LoRA adapters, or right before the
  generation point for aLoRA adapters.
- **KV-cache hiding.** Control tokens use group-based control dimensions
  on Q/K so adapters never see another adapter's internal cache state.
  This is what lets twelve adapters coexist without crosstalk.
- **Frozen during inference.** Adapters are pre-trained separately and
  baked into the checkpoint. There's no per-request adapter loading.
- **One forward pass.** Adding adapters costs a small fraction of the
  base model's parameters, and adapter computation reuses the base
  forward — there's no second model to load.

### Why this matters

Deploying twelve specialized models traditionally means twelve
checkpoints, twelve servers, twelve GPU allocations. Granite Switch
collapses that to one checkpoint that picks the right specialist per
request, addressed by a single token in the prompt.

### Learn more

- [Model card](https://huggingface.co/ibm-granite/granite-switch-4.1-3b-preview)
  — full architecture description, supported adapters, and benchmarks.
- [Source on GitHub](https://github.com/generative-computing/granite-switch)
  — composer (build your own switch from base + adapters), HF backend,
  vLLM backend.
- [PyPI: `granite-switch`](https://pypi.org/project/granite-switch/) —
  `pip install "granite-switch[hf]"` for the HF backend.
"""

with gr.Blocks(title="Granite Switch 4.1 3B", theme=gr.themes.Soft()) as demo:
    gr.Markdown(INTRO)

    with gr.Tabs():
        with gr.Tab("Demo"):
            with gr.Row():
                with gr.Column(scale=2):
                    adapter = gr.Dropdown(
                        choices=adapter_choices_grouped(),
                        value="(none — base model)",
                        label="Adapter",
                    )
                with gr.Column(scale=1):
                    max_tokens = gr.Slider(
                        minimum=8, maximum=512, value=128, step=8,
                        label="Max new tokens",
                    )

            description = gr.Markdown(
                "**Base** — the unmodified Granite 4.1 3B model.",
            )

            with gr.Row():
                with gr.Column():
                    messages = gr.Code(
                        label="Messages (JSON list of {role, content})",
                        language="json", lines=14,
                    )
                with gr.Column():
                    documents = gr.Code(
                        label="Documents (JSON list, optional)",
                        language="json", lines=14,
                    )

            with gr.Row():
                run_btn = gr.Button("Generate", variant="primary", scale=2)
                show_prompt_btn = gr.Button("Show rendered prompt", scale=1)

            output = gr.Textbox(label="Adapter output", lines=10)
            status = gr.Markdown("", elem_id="gen-status")

            with gr.Accordion("Rendered prompt (with control token)", open=False):
                rendered = gr.Textbox(
                    label="", lines=14,
                    info="The exact string passed to the model. Look for the "
                         "adapter control token — that's the switch firing.",
                )

            adapter.change(
                load_example,
                inputs=[adapter],
                outputs=[messages, documents, max_tokens, description],
            )
            run_btn.click(
                generate,
                inputs=[adapter, messages, documents, max_tokens],
                outputs=[output, status],
            )
            show_prompt_btn.click(
                render_prompt,
                inputs=[adapter, messages, documents],
                outputs=rendered,
            )

            demo.load(
                load_example,
                inputs=[adapter],
                outputs=[messages, documents, max_tokens, description],
            )

        with gr.Tab("About the switch"):
            gr.Markdown(ABOUT)


if __name__ == "__main__":
    demo.launch()