Function generate

trainer/src/inference_server/model_benchmark.py:103–151

(request: Input)

)
@modal.web_endpoint(method="POST", label=APP_NAME)
async def generate(request: Input) -> Output:
    logging.info(f"Generating for model {request.model}")

    sample_params = SamplingParams(
        n=request.n,
        temperature=request.temperature,
        max_tokens=request.max_tokens,
    )

    request_id = random_uuid()

    logging.info(f"Generating for request {request_id}")

    lora_request = None
    if TESTING_LORA:
        lora_request = LoRARequest(
            request.model, abs(hash(request.model)), lora_model_cache_dir(MODEL_ID)
        )

    output_generator = engine.generate(
        request.prompt,
        sample_params,
        request_id=request_id,
        lora_request=lora_request,
    )

    final_output: Union[RequestOutput, None] = None
    async for request_output in output_generator:
        final_output = request_output

    if final_output is None:
        raise Exception("No output generated")

    prompt_tokens = len(final_output.prompt_token_ids)
    completion_tokens = sum(len(x.token_ids) for x in final_output.outputs)

    output = Output(
        id=request_id,
        choices=[
            Choice(text=choice.text, finish_reason=choice.finish_reason)
            for choice in final_output.outputs
        ],
        usage=Usage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
        ),
    )

    return output