from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import time
torch.manual_seed(0)

#path = "output/AdvertiseGenLoRA_lora_finetune/xxx/checkpoint-3000" # xxx：系统时间路径
path = 'checkpoint/miniCPM-bf16'
tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16, device_map='cuda', trust_remote_code=True)

start_time = time.time()
responds, history = model.chat(tokenizer, "山东省最高的山是哪座山, 它比黄山高还是矮？差距多少？", temperature=0.5, top_p=0.8, repetition_penalty=1.02)
print("infer time：", time.time() - start_time, "s")
print(responds)