Function generate

trainer/src/lora_inference/experiments.py:89–148

(self, request)

    image=vllm_image,
)
async def generate(self, request):
    logging.info(f"Processing for model {request.lora_model}")
    lora_dir = lora_model_cache_dir(request.lora_model)

    if request.lora_model not in self.seen_models:
        if not os.path.exists(lora_model_cache_dir(request.lora_model)):
            logging.info(f"Couldn't find model, reloading {lora_dir}")
            stub.volume.reload()

        if not os.path.exists(lora_model_cache_dir(request.lora_model)):
            raise Exception(f"Couldn't find model {lora_dir}after reloading!")

        self.seen_models.add(request.lora_model)

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

    lora_request = LoRARequest(
        request.lora_model,
        abs(hash(request.lora_model)),
        lora_dir,
    )

    request_id = random_uuid()

    logging.info(f"Generating for request {request_id}")
    output_generator = self.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:
        # TODO: support streaming
        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,
    #     ),
    # )