from  transformers import AutoTokenizer, AutoModelForCausalLM


if __name__ == '__main__':
    model_name = "baidu/ERNIE-4.5-0.3B-PT"
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)

    prompt = "Give me a short introduction to large language model."
    messages = [
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([text], add_special_tokens=True, return_tensors="pt")

    # conduct text completion
    generated_ids = model.generate(
        model_inputs.input_ids,
        max_new_tokens=1024
    )
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

    # decode the generated ids
    generate_text = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
    print("generate_text:", generate_text)