# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.13.8
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig

# Mengunduh tokenizer dari Hugging Face
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")

# Mengunduh dan memuat model AutoModelForCausalLM dalam mode 8-bit dan secara otomatis memetakan ke perangkat yang tersedia
model = AutoModelForCausalLM.from_pretrained(
    "microsoft/phi-2",
    load_in_8bit=True,      # Menggunakan mode 8-bit untuk mengurangi penggunaan memori
    device_map="auto"       # Secara otomatis memetakan model ke perangkat yang tersedia (CPU/GPU)
)

# Menggunakan model PEFT yang telah dilatih sebelumnya
model = PeftModel.from_pretrained(model, "petualang/Phi2-Doctor20")

# Anda dapat menambahkan kode di bawah ini untuk memverifikasi bahwa model dan tokenizer telah dimuat dengan benar
# atau untuk menguji prediksi sederhana.

# Contoh input teks
input_text = "Hallo Doctor. I have a fever now, what should i do?"

# Tokenisasi input
input_ids = tokenizer.encode(input_text, return_tensors="pt").to(model.device)

# Menghasilkan teks dengan model
output_ids = model.generate(input_ids, max_length=50)

# Dekode hasil prediksi menjadi teks
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)

print(output_text)