Spaces:
Sleeping
Sleeping
File size: 7,759 Bytes
be82719 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | """Generate view (miru's Logging Mode): bulk generation with probability logging.
Single-user server-side run state (tracer + originating inputs) lives in
``state.get_logging_session()`` — the streaming handler updates it so Stop and
Continue can reach the live tracer, and Continue is invalidated when the model
changes (object identity, like upstream).
"""
from __future__ import annotations
import time
from collections.abc import Iterator
from api.helpers import ChatValidationError, ui_sampling_params
from api.interactive import _reset_tracer_for_mode
from api.models import model_manager
from api.serialize import error_payload, fig_json
from api.state import get_logging_session
from miru_tracer.core.logging_config import get_logger
from miru_tracer.core.tracer import LLMTracer
from miru_tracer.visualization.plots import (
get_generation_stats,
plot_probability_visualizations,
)
logger = get_logger(__name__)
def _finalize(tracer, params, heatmap_ranks, prob_mode) -> dict:
"""Stats, plots, and export for a finished (or stopped) run."""
stats = get_generation_stats(tracer.history)
ranks = int(heatmap_ranks) if heatmap_ranks else 10
if tracer.history:
ranks = min(ranks, len(tracer.history[0].top_k_tokens))
figures = plot_probability_visualizations(
tracer.history,
top_k=ranks,
probability_mode=prob_mode,
temperature=params.temperature,
)
return {
"stats": stats,
"fig_heatmap": fig_json(figures[0]) if figures else None,
"fig_confidence": fig_json(figures[1]) if len(figures) > 1 else None,
"export": tracer.export_to_dict(params),
}
def _stream_run(
tracer, params, originals, max_new_tokens, log_top_k, log_full_probs,
stop_at_eos, heatmap_ranks, prob_mode, start_time,
) -> Iterator[dict]:
"""Shared streaming loop for generate and continue."""
tokens_before = len(tracer.history)
for step_data in tracer.generate_stream(
max_new_tokens=int(max_new_tokens),
params=params,
log_top_k=max(int(log_top_k or 10), 1),
log_full_probs=bool(log_full_probs),
stop_at_eos=bool(stop_at_eos),
):
yield {
"ok": True,
"type": "progress",
"text": tracer.get_generated_text(),
"progress": {
"step": len(tracer.history),
"new_tokens": len(tracer.history) - tokens_before,
"last_token": step_data.token_text,
},
}
final = _finalize(tracer, params, heatmap_ranks, prob_mode)
logger.info(
f"Generation finished: {len(tracer.history)} tokens in "
f"{time.time() - start_time:.2f}s"
)
yield {
"ok": True,
"type": "final",
"text": tracer.get_generated_text(),
"originals": list(originals),
**final,
}
def logging_generate(
mode: str, prompt: str, chat_json: str, raw_text: str,
thinking: str, think_prefill: str,
max_new_tokens: int, strategy: str, temperature: float,
top_k: int, top_p: float,
stop_at_eos: bool, log_top_k: int, heatmap_ranks: int,
log_full_probs: bool, prob_mode: str,
) -> Iterator[dict]:
model = model_manager.get_model()
tokenizer = model_manager.get_tokenizer()
device = model_manager.get_device()
run = get_logging_session()
if model is None or tokenizer is None:
yield error_payload("No model loaded. Load one in the Model view first.")
return
if max_new_tokens is None or max_new_tokens < 1:
yield error_payload("Maximum new tokens must be at least 1")
return
start_time = time.time()
try:
params = ui_sampling_params(strategy, temperature, top_k, top_p)
tracer = LLMTracer(model, tokenizer, device)
_reset_tracer_for_mode(
tracer, mode, prompt, chat_json, raw_text, thinking, think_prefill
)
originals = (mode, prompt, chat_json, raw_text, thinking, think_prefill)
with run.lock:
run.tracer = tracer # reachable by Stop mid-run
run.originals = originals
logger.info(
f"Logging generation started: mode={mode}, "
f"max_tokens={max_new_tokens}, strategy={strategy}"
)
yield from _stream_run(
tracer, params, originals, max_new_tokens, log_top_k,
log_full_probs, stop_at_eos, heatmap_ranks, prob_mode, start_time,
)
except ChatValidationError as e:
with run.lock:
run.tracer = None
run.originals = None
yield error_payload(str(e))
except Exception as e:
logger.error(f"Generation error: {e}", exc_info=True)
with run.lock:
run.tracer = None
run.originals = None
yield error_payload(f"Error during generation:\n\n{e}", trace=True)
def logging_continue(
max_new_tokens: int, strategy: str, temperature: float,
top_k: int, top_p: float,
stop_at_eos: bool, log_top_k: int, heatmap_ranks: int,
log_full_probs: bool, prob_mode: str,
) -> Iterator[dict]:
"""Resume the same tracer; history accumulates across runs."""
run = get_logging_session()
with run.lock:
tracer = run.tracer
originals = run.originals
if tracer is None:
yield error_payload("No previous generation to continue from.")
return
if max_new_tokens is None or max_new_tokens < 1:
yield error_payload("Maximum new tokens must be at least 1")
return
if model_manager.get_model() is None or tracer.model is not model_manager.get_model():
yield error_payload(
"Model has been unloaded or changed.\n\nPlease start a new generation."
)
return
start_time = time.time()
try:
params = ui_sampling_params(strategy, temperature, top_k, top_p)
logger.info(f"Continuing generation: max_tokens={max_new_tokens}")
yield from _stream_run(
tracer, params, originals, max_new_tokens, log_top_k,
log_full_probs, stop_at_eos, heatmap_ranks, prob_mode, start_time,
)
except Exception as e:
logger.error(f"Continuation error: {e}", exc_info=True)
yield error_payload(f"Error during continuation:\n\n{e}", trace=True)
def logging_stop() -> dict:
"""Cooperative stop; the running stream finalizes the partial run itself."""
run = get_logging_session()
with run.lock:
tracer = run.tracer
if tracer is None:
return error_payload("No generation running.")
tracer.request_stop()
logger.info(f"Stop requested after {len(tracer.history)} tokens")
return {"ok": True}
def logging_replot(heatmap_ranks: int, prob_mode: str, temperature: float) -> dict:
"""Re-render plots with new display settings — no regeneration."""
run = get_logging_session()
with run.lock:
tracer = run.tracer
if tracer is None or not tracer.history:
return {"ok": True, "fig_heatmap": None, "fig_confidence": None}
try:
ranks = min(
int(heatmap_ranks) if heatmap_ranks else 10,
len(tracer.history[0].top_k_tokens),
)
figures = plot_probability_visualizations(
tracer.history,
top_k=ranks,
probability_mode=prob_mode,
temperature=float(temperature),
)
return {
"ok": True,
"fig_heatmap": fig_json(figures[0]) if figures else None,
"fig_confidence": fig_json(figures[1]) if len(figures) > 1 else None,
}
except Exception as e:
logger.error(f"Error refreshing visualizations: {e}", exc_info=True)
return error_payload(str(e))
|