Open goog opened 1 year ago
from transformers import AutoModelForCausalLM, AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("baichuan-inc/Baichuan-7B", trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained("baichuan-inc/Baichuan-7B", device_map="auto", trust_remote_code=True) inputs = tokenizer('犹有渔人数家住', return_tensors='pt') inputs = inputs.to('cuda:0') pred = model.generate(**inputs, max_new_tokens=128,repetition_penalty=1.1, do_sample=True, top_k=30, top_p=0.92, #temperature=0.001, num_return_sequences=5) #print(pred.cpu()) sample_outputs = pred.cpu() for i, sample_output in enumerate(sample_outputs): print("{}: {}".format(i, tokenizer.decode(sample_output, skip_special_tokens=True))) #print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))
Required prerequisites
Questions
Checklist