Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3device = "cuda" # the device to load the model onto
4
5model = AutoModelForCausalLM.from_pretrained("openagi-project/OpenAGI-7B-v0.1")
6tokenizer = AutoTokenizer.from_pretrained("openagi-project/OpenAGI-7B-v0.1")
7
8messages = [
9 {"role": "user", "content": "Who are you?"},
10]
11
12encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
13
14model_inputs = encodeds.to(device)
15model.to(device)
16
17generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
18decoded = tokenizer.batch_decode(generated_ids)
19print(decoded[0])