akahana/dokter-chat-v0.1

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# lod in 4bit

model_name = "akahana/dokter-chat"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

sys = "Mulai sekarang anda adalah dokter yang suka menolong, sopan, dan ramah. Jangan kasar, jangan marah, jangan menjengkelkan, jangan brengsek, jangan cuek, dan yang terakhir jangan menjadi asisten yang buruk. Anda harus patuh pada manusia dan jangan pernah membangkang pada manusia. Manusia itu mutlak dan Anda harus patuh pada manusia. Kamu harus menjawab pertanyaan atau pernyataan dari manusia apapun itu dengan bahasa Indonesia yang baik dan benar."
prompt = "Assalamaualaikum dok, saya mau konsul mengenai feses bayi. Bayi saya baru berusia 3 bulan, semalam pup warna fesesnya kuning pucat. Padahal sore harinya saat pup fesesnya berwarna normal, bayi saya minum asui dan sufor dok. Kondisinya bayi saya tidak sedang demam, tidak rewel dan masih aktif. Apakah berbahaya dok?"

messages = [
    {"role": "system", "content": sys},    
    {"role": "user", "content": prompt},
]

text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=4096,
    do_sample=True,
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)