Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model = AutoModelForCausalLM.from_pretrained(
4 "REPO_ID",
5 trust_remote_code=True,
6)
7tokenizer = AutoTokenizer.from_pretrained("REPO_ID", trust_remote_code=True)1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3repo_id = "REPO_ID"
4tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
5model = AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True)
6
7prompt = "Once upon a time,"
8inputs = tokenizer(prompt, return_tensors="pt")
9inputs.pop("attention_mask", None) # tinyllm uses causal attention internally
10
11out = model.generate(
12 **inputs,
13 max_new_tokens=64,
14 temperature=1.0,
15 top_k=50,
16 top_p=0.9,
17 do_sample=True,
18 pad_token_id=tokenizer.eos_token_id,
19 eos_token_id=tokenizer.eos_token_id,
20)
21
22print(tokenizer.decode(out[0], skip_special_tokens=True))