Repository: LAION-AI/Open-Assistant
File: inference/worker/work.py (Python)

Code before:
    label = label.replace("<pad>", "").strip()

    if "caution" in label and safety_level > 1:
        return prepare_safe_prompt(prompt, label, rots)
    elif "intervention" in label and safety_level > 0:
        return prepare_safe_prompt(prompt, label, rots)
    else:
        return prompt


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, settings.safety_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(),

Review comment (suggestion):
We are already receiving a desired safety level from the backend (which in turn comes from the frontend as a user setting) as part of the `SafetyParameters` in the work request. So we can remove the `safety_level` setting from `settings.py` and this line can become:

`prompt = get_safety_opinion(prompt, safety_response.outputs, work_request.safety_parameters.level)`

Code after:
    label = label.replace("<pad>", "").strip()

    if "caution" in label and safety_level > 1:
        return prepare_safe_prompt(prompt, label, rots)
    elif "intervention" in label and safety_level > 0:
        return prepare_safe_prompt(prompt, label, rots)
    else:
        return prompt


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(),