from transformers import AutoTokenizer, AquilaForCausalLM
import torch
from cyg_conversation import default_conversation, covert_prompt_to_input_ids_with_history

tokenizer = AutoTokenizer.from_pretrained("BAAI/AquilaChat-7B")
model = AquilaForCausalLM.from_pretrained("BAAI/AquilaChat-7B")
model.eval()
model.to("cuda:4")
vocab = tokenizer.vocab
print(len(vocab))

text = "请给出10个要到北京旅游的理由。"

tokens = covert_prompt_to_input_ids_with_history(text, history=[], tokenizer=tokenizer, max_token=512)

tokens = torch.tensor(tokens)[None,].to("cuda:4")

out = model.generate(tokens, do_sample=True, max_length=512, eos_token_id=100007)[0]
out = tokenizer.decode(out.cpu().numpy().tolist())

print(out)