Function event_generator

project/monitor/monitor/server.py:214–257

()

    nv.set_neuron_interventions(interventions)

    def event_generator():
        gen = cast(
            Generator[int | GenerateOutput, None, None],
            nv.send_message(
                SUBJECT,
                message,
                max_new_tokens=max_new_tokens or 64,
                temperature=temperature or 1.0,
                stream=True,
            ),
        )
        chat_tokens: list[ChatToken] = []

        for update in gen:
            if isinstance(update, int):
                chat_tokens.append(ChatToken(token=SUBJECT.decode(update)))
                data = json.dumps([ct.model_dump() for ct in chat_tokens])
                yield f"data: {data}\n\n"
            else:
                # Collate tokenwise log probs
                index_to_log_probs: dict[int, list[tuple[str, float]]] = {}
                tokenwise_log_probs = update.tokenwise_log_probs
                for i, (token_ids, log_probs) in enumerate(tokenwise_log_probs):
                    cur_log_probs = [
                        (SUBJECT.decode(token_id), log_prob)
                        for token_id, log_prob in zip(token_ids[0], log_probs[0])
                    ]
                    # output_ids_BT.shape[1] is the total sequence length, len(tokenwise_log_probs) is the number of tokens generated
                    index_to_log_probs[
                        i + update.output_ids_BT.shape[1] - len(tokenwise_log_probs)
                    ] = cur_log_probs

                # Update chat tokens with top log probs
                chat_tokens = [
                    ChatToken(**(ct.model_dump() | {"top_log_probs": index_to_log_probs.get(i)}))
                    for i, ct in enumerate(chat_tokens)
                ]
                data = json.dumps([ct.model_dump() for ct in chat_tokens])

                yield f"data: {data}\n\n"
                yield "data: [DONE]\n\n"

        nv.clear_neuron_interventions()
        SESSIONS[session_id] = (ChatConversation(**nv.model_input.model_dump()), nv.filter)

    return StreamingResponse(event_generator(), media_type="text/event-stream")