|
|
|
|
|
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) |
|
|
