# Please use transformers==4.45.0

import torch
import transformers

model_id = "GoToCompany/gemma2-9b-cpt-sahabatai-v1-instruct"

pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
)

terminators = [
    pipeline.tokenizer.eos_token_id,
    pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
]

# Javanese
messages = [
   {"role": "user", "content": "Sopo wae sing ana ing Punakawan?"}
]

outputs = pipeline(
    messages,
    max_new_tokens=256,
    eos_token_id=terminators,
)
print(outputs[0]["generated_text"][-1])


# Sundanese
messages = [
    {"role": "user", "content": "Kumaha caritana si Kabayan?"},
]

outputs = pipeline(
    messages,
    max_new_tokens=256,
    eos_token_id=terminators,
)
print(outputs[0]["generated_text"][-1])