src/espelho/ui.py

"""UI Gradio do stand: chat à esquerda, painel do lens à direita.

Uso: uv run python -m espelho.ui
"""

import html
import queue
import threading
import time

import gradio as gr

try:  # ZeroGPU: decorator é efeito-nulo fora do Space (verificado nos docs)
    import spaces
    _ZERO_GPU = bool(__import__("os").environ.get("SPACES_ZERO_GPU"))
    _gpu = spaces.GPU(duration=120)
except ImportError:  # ambiente local sem o pacote
    _ZERO_GPU = False
    _gpu = None

from espelho import config
from espelho.filters import (
    aggregate_chips, baseline_blacklist, heatmap_grid, load_stopwords,
)
from espelho.model import (
    available_models, generate_with_trace, load_model, on_space, text_config,
)

RODAPE = (
    "Demonstração de pesquisa — não digite dados pessoais. "
    "As conversas não são armazenadas."
)

NOMES = {
    "google/gemma-2-2b-it": "Gemma 2 (2B)",
    "google/gemma-3-4b-it": "Gemma 3 (4B)",
    "google/gemma-2-9b-it": "Gemma 2 (9B)",
    "google/gemma-3-12b-it": "Gemma 3 (12B)",
    "Qwen/Qwen3-8B": "Qwen3 (8B)",
}

PLACEHOLDER = "Digite aqui e pressione Enter…"

CSS = """
.chatcol .message, .chatcol .message p { font-size: 20px !important; }
.chatcol textarea { font-size: 20px !important; }
#chips .chip {
  display: inline-block; font-size: 18px; padding: 6px 14px; margin: 4px;
  border-radius: 16px; background: #e8eaf6; border: 1px solid #9fa8da;
}
#chips .vazio { font-size: 18px; color: #666; }
#rodape {
  position: fixed; bottom: 0; left: 0; right: 0; z-index: 1000;
  text-align: center; padding: 6px; font-size: 16px;
  background: #fff3cd; border-top: 1px solid #ccc;
}
#heatmap table { border-collapse: collapse; font-size: 14px; }
#heatmap td, #heatmap th { border: 1px solid #ddd; padding: 3px 6px; text-align: center; }
footer { display: none !important; }
"""

# A Source Sans Pro empacotada no tema padrão renderiza "s" como "ſ" (s longo)
# no macOS; fontes do sistema evitam o bug e qualquer fetch remoto no stand.
TEMA = gr.themes.Default(
    font=["-apple-system", "system-ui", "Helvetica Neue", "Arial", "sans-serif"],
    font_mono=["ui-monospace", "SFMono-Regular", "Menlo", "monospace"],
)


def _chips_html(chips: list[dict]) -> str:
    if not chips:
        return "<p class='vazio'>Nada destacado neste turno.</p>"
    spans = [
        f"<span class='chip' title='melhor rank {c['rank'] + 1} na camada {c['layer']}'>"
        f"{html.escape(c['token'])}</span>"
        for c in chips
    ]
    return "<div>" + "".join(spans) + "</div>"


def _heatmap_html(grid: dict[str, dict[int, int]], band: list[int]) -> str:
    if not grid:
        return ""
    head = "".join(f"<th>{l}</th>" for l in band)
    rows = []
    for token, cells in grid.items():
        tds = []
        for l in band:
            rank = cells.get(l)
            if rank is None:
                tds.append("<td style='background:#f5f5f5'></td>")
            else:
                # rank 0 (forte) -> verde escuro; rank 24 (fraco) -> quase branco
                alpha = max(0.15, 1.0 - rank / config.LENS_TOP_K)
                tds.append(
                    f"<td style='background:rgba(46,125,50,{alpha:.2f})'>{rank + 1}</td>"
                )
        rows.append(f"<tr><th>{html.escape(token)}</th>{''.join(tds)}</tr>")
    return (
        "<table><tr><th>token \\ camada</th>" + head + "</tr>"
        + "".join(rows) + "</table>"
    )


def _explicacao(model) -> str:
    num_camadas = text_config(model).num_hidden_layers
    banda = config.layer_band(num_camadas)
    return (
        f"Este modelo processa o texto em {num_camadas} camadas. "
        f"Com a técnica **logit lens**, traduzimos em palavras o estado "
        f"interno das camadas {banda[0]} a {banda[-1]} enquanto ele "
        "escrevia. Os conceitos abaixo estavam ativados nesse processo, "
        "mas ficaram de fora da resposta final."
    )


def build_app(bundles: dict[str, list], ativo_id: str | None = None):
    """Monta a UI sobre bundles [tokenizer, model, blacklist|None] por modelo.

    No ZeroGPU todos os modelos do seletor chegam carregados no boot (a
    emulação CUDA não intercepta cargas em runtime) e a calibração da
    blacklist fica para a 1ª mensagem de cada modelo (roda com GPU real).
    No local, a calibração pendente roda aqui e outros modelos cacheados
    podem ser carregados sob demanda na troca.
    """
    stopwords = load_stopwords()
    if not _ZERO_GPU:
        for b in bundles.values():
            if b[2] is None:
                print("Calibrando lista-negra do lens (baseline)...")
                b[2] = baseline_blacklist(b[0], b[1])
    ativo = {"id": ativo_id or next(iter(bundles))}
    if on_space():
        disponiveis = list(bundles)
    else:
        disponiveis = available_models() or list(bundles)
    ultima_atividade = {"t": time.time()}

    def responder(mensagem: str, historico: list[dict]):
        ultima_atividade["t"] = time.time()
        bundle = bundles[ativo["id"]]
        if bundle[2] is None:  # calibração adiada (ZeroGPU): 1ª chamada
            bundle[2] = baseline_blacklist(bundle[0], bundle[1])
        tok, mdl, blacklist = bundle
        historico = list(historico or [])
        mensagem = (mensagem or "").strip()
        if not mensagem:
            yield historico, gr.skip(), gr.skip(), ""
            return

        historico.append({"role": "user", "content": mensagem})
        historico.append({"role": "assistant", "content": ""})

        fila: queue.Queue = queue.Queue()
        resultado: dict = {}

        def gerar():
            try:
                _, trace = generate_with_trace(
                    tok, mdl, historico[:-1],
                    max_new_tokens=config.MAX_NEW_TOKENS,
                    temperature=config.TEMPERATURE,
                    on_text=fila.put,
                )
                resultado["trace"] = trace
            except Exception as exc:  # robustez de stand: erro vira mensagem
                resultado["erro"] = str(exc)
            finally:
                fila.put(None)

        threading.Thread(target=gerar, daemon=True).start()

        while True:
            pedaco = fila.get()
            if pedaco is None:
                break
            historico[-1]["content"] += pedaco
            ultima_atividade["t"] = time.time()
            yield historico, gr.skip(), gr.skip(), ""

        trace = resultado.get("trace")
        if trace is None:
            historico[-1]["content"] = (
                "Desculpe, algo deu errado. Toque em Recomeçar e tente de novo."
            )
            yield historico, _chips_html([]), "", ""
            return

        chips = aggregate_chips(
            trace, mensagem,
            historico[-1]["content"], stopwords, blacklist=blacklist,
        )
        heat = _heatmap_html(heatmap_grid(trace, chips), trace.band_layers)
        ultima_atividade["t"] = time.time()
        yield historico, _chips_html(chips), heat, ""

    if _gpu is not None:
        responder = _gpu(responder)

    def recomecar():
        ultima_atividade["t"] = time.time()
        return [], "", "", ""

    def tique(historico):
        if historico and time.time() - ultima_atividade["t"] > config.IDLE_RESET_SECONDS:
            ultima_atividade["t"] = time.time()
            return [], "", "", ""
        return gr.skip(), gr.skip(), gr.skip(), gr.skip()

    def trocar_modelo(nome: str):
        """Saídas: chat, chips, heatmap, caixa, explicação, seletor.

        Durante a carga, caixa e seletor ficam desabilitados — nada digitado
        se perde e não há troca dupla no meio da carga.
        """
        ultima_atividade["t"] = time.time()
        model_id = next((k for k, v in NOMES.items() if v == nome), nome)
        if model_id == ativo["id"]:
            yield (gr.skip(),) * 6
            return
        if model_id not in bundles:
            yield (
                [], "", "",
                gr.update(value="", interactive=False,
                          placeholder="Aguarde: carregando o modelo…"),
                f"Carregando **{nome}** — pode levar alguns instantes…",
                gr.update(interactive=False),
            )
            try:
                tok, mdl = load_model(model_id)
            except Exception as exc:
                yield (
                    gr.skip(), gr.skip(), gr.skip(),
                    gr.update(interactive=True, placeholder=PLACEHOLDER),
                    f"Não consegui carregar {nome} ({exc}). Voltando ao anterior.",
                    gr.update(interactive=True,
                              value=NOMES.get(ativo["id"], ativo["id"])),
                )
                return
            bundles[model_id] = [
                tok, mdl, None if _ZERO_GPU else baseline_blacklist(tok, mdl)
            ]
        ativo["id"] = model_id
        # Troca de modelo zera a conversa: traces de modelos diferentes não se misturam.
        yield (
            [], "", "",
            gr.update(value="", interactive=True, placeholder=PLACEHOLDER),
            _explicacao(bundles[model_id][1]),
            gr.update(interactive=True),
        )

    with gr.Blocks(title="Padrões de ativação") as app:
        with gr.Row():
            gr.Markdown("# Padrões de ativação")
            seletor = gr.Dropdown(
                choices=[NOMES.get(m, m) for m in disponiveis],
                value=NOMES.get(ativo["id"], ativo["id"]),
                label="Modelo",
                interactive=len(disponiveis) > 1,
                scale=0, min_width=220,
            )
        with gr.Row():
            with gr.Column(scale=3, elem_classes=["chatcol"]):
                chat = gr.Chatbot(label="Conversa", height=480)
                caixa = gr.Textbox(
                    label="Sua mensagem",
                    placeholder=PLACEHOLDER,
                    submit_btn=True,
                )
                with gr.Row():
                    botao_reset = gr.Button(
                        "Recomeçar", variant="secondary", size="sm", scale=0,
                    )
            with gr.Column(scale=2):
                gr.Markdown("## Conceitos ativados dentro do modelo que não apareceram na resposta")
                explicacao = gr.Markdown(
                    _explicacao(bundles[ativo["id"]][1]), elem_id="explicacao"
                )
                chips_html = gr.HTML(elem_id="chips")
                with gr.Accordion("ver por camada", open=False):
                    heat_html = gr.HTML(elem_id="heatmap")
        gr.HTML(f"<div id='rodape'>{RODAPE}</div>")

        saidas = [chat, chips_html, heat_html, caixa]
        caixa.submit(responder, [caixa, chat], saidas, api_name="enviar")
        botao_reset.click(recomecar, [], saidas, api_name="recomecar")
        gr.Timer(5).tick(tique, [chat], saidas)
        seletor.change(
            trocar_modelo, [seletor], saidas + [explicacao, seletor],
            api_name="modelo",
        )

    return app


def main() -> None:
    print(f"Carregando {config.ACTIVE_MODEL} ...")
    tokenizer, model = load_model()
    app = build_app({config.ACTIVE_MODEL: [tokenizer, model, None]})
    # Gradio 6: css e theme são parâmetros do launch(), não do Blocks.
    app.launch(
        server_name="127.0.0.1", server_port=7860,
        show_error=True, css=CSS, theme=TEMA,
    )


if __name__ == "__main__":
    main()