Function handle_work_request

inference/worker/work.py:103–187

(
    ws: websocket.WebSocket,
    tokenizer: transformers.PreTrainedTokenizer,
    work_request: inference.WorkRequest,
    worker_config: inference.WorkerConfig,
)

def handle_work_request(
    ws: websocket.WebSocket,
    tokenizer: transformers.PreTrainedTokenizer,
    work_request: inference.WorkRequest,
    worker_config: inference.WorkerConfig,
):
    prompt, parameters = make_prompt_and_parameters(tokenizer=tokenizer, work_request=work_request)
    logger.debug(f"Prompt: {prompt}")

    model_config = worker_config.model_config

    # Only send safety request if work request safety level is not 0
    if settings.enable_safety and work_request.safety_parameters.level:
        safety_request = inference.SafetyRequest(inputs=prompt, parameters=work_request.safety_parameters)
        safety_response = get_safety_server_response(safety_request)
        prompt = get_safety_opinion(prompt, safety_response.outputs, work_request.safety_parameters.level)
        logger.debug(f"Safe prompt: {prompt}")

    stream_response = None
    token_buffer = utils.TokenBuffer(stop_sequences=parameters.stop)
    if model_config.is_lorem:
        stream_events = utils.lorem_events(parameters.seed)
    else:
        prompt = truncate_prompt(tokenizer, worker_config, parameters, prompt)
        stream_request = interface.GenerateStreamRequest(
            inputs=prompt,
            parameters=parameters,
        )
        stream_events = get_inference_server_stream_events(stream_request)

    generated_ids = []
    decoded_text = ""
    for stream_response in stream_events:
        if stream_response.is_error:
            logger.error(f"Error from inference server: {stream_response.error}")
            utils.send_response(
                ws,
                inference.ErrorResponse(
                    request_id=work_request.id,
                    error=stream_response.error,
                    metrics=inference.WorkerMetricsInfo(),
                ),
            )
            raise RuntimeError(f"Error from inference server: {stream_response.error}")
        token = stream_response.token

        if model_config.is_llama:
            generated_ids.append(token.id)
            try:
                with tokenizer_lock:
                    text = tokenizer.decode(generated_ids, skip_special_tokens=True)
                new_text = text[len(decoded_text) :]
                if not decoded_text:
                    new_text = new_text.lstrip()
            except Exception:
                text = decoded_text
                new_text = ""
            token.text = new_text