1import torch
2from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
3
4# IMPORTANT: This model requires a modern GPU (Ampere/Ada/Hopper) that supports bfloat16.
5# Running on CPU or older GPUs (like T4) will result in NaN generation errors.
6device = "cuda" if torch.cuda.is_available() else "cpu"
7
8repo_id = "neosantara/wader-100m"
9tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
10model = AutoModelForCausalLM.from_pretrained(
11 repo_id,
12 trust_remote_code=True,
13 torch_dtype=torch.bfloat16
14).to(device)
15
16# Chat
17messages = [{"role": "user", "content": "Halo, apa kabar?"}]
18prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
19input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
20
21outputs = model.generate(
22 input_ids,
23 max_new_tokens=200,
24 temperature=0.7,
25 top_p=0.9,
26 pad_token_id=tokenizer.eos_token_id,
27 do_sample=True
28)
29print(tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True))