import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

adapter_id = "Reubencf/gemma-3-4b-it-amharic-lora"
base_id = "NotoriousH2/gemma-3-4b-it-TextOnly"
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32

tokenizer = AutoTokenizer.from_pretrained(adapter_id, use_fast=True)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

base_model = AutoModelForCausalLM.from_pretrained(
    base_id,
    torch_dtype=dtype,
    device_map={"": device},
)
model = PeftModel.from_pretrained(base_model, adapter_id, torch_device=device)
model.eval()

messages = [
    {"role": "user", "content": "Write a short greeting in Amharic."}
]
inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_tensors="pt",
).to(device)

with torch.inference_mode():
    output = model.generate(
        inputs,
        max_new_tokens=128,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        pad_token_id=tokenizer.pad_token_id,
    )

response = output[0][inputs.shape[-1]:]
print(tokenizer.decode(response, skip_special_tokens=True))