Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4torch.set_default_device("cuda")
5
6model_path = "olliai/Maika-Buddy-2B7-Chat-v2"
7
8model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto", trust_remote_code=True)
9tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
10
11prompt = """A chat between a curious kid and an artificial intelligence assistant. The assistant gives helpful, easy-to-understand, detailed, and polite answers to the questions of the kid.
12Kid: I have 10 apples. I gave Bob 2 pencil. How many apples do I have now?
13AI: """
14
15inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False)
16
17outputs = model.generate(**inputs, max_length=200,
18 do_sample=True,
19 repetition_penalty=1.1,
20 temperature=0.3,
21 top_k=50,
22 top_p=0.95,
23 eos_token_id=tokenizer.encode("<|endoftext|>"))
24text = tokenizer.batch_decode(outputs)[0]
25print(text)
26