Function generate

applications/Colossal-LLaMA/inference/inference_example.py:26–51

(args)

@torch.inference_mode()
def generate(args):
    model, tokenizer = load_model(model_path=args.model_path, device=args.device)

    if args.prompt_style == "sft":
        conversation = default_conversation.copy()
        conversation.append_message("Human", args.input_txt)
        conversation.append_message("Assistant", None)
        input_txt = conversation.get_prompt()
    else:
        BASE_INFERENCE_SUFFIX = "\n\n->\n\n"
        input_txt = f"{args.input_txt}{BASE_INFERENCE_SUFFIX}"

    inputs = tokenizer(input_txt, return_tensors="pt").to(args.device)
    num_input_tokens = inputs["input_ids"].shape[-1]
    output = model.generate(
        **inputs,
        max_new_tokens=args.max_new_tokens,
        do_sample=args.do_sample,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
        num_return_sequences=1,
    )
    response = tokenizer.decode(output.cpu()[0, num_input_tokens:], skip_special_tokens=True)
    logger.info(f"\nHuman: {args.input_txt} \n\nAssistant: \n{response}")
    return response


if __name__ == "__main__":