import gradio as gr                          # Mengimpor pustaka Gradio untuk membuat UI
from transformers import AutoTokenizer, AutoModelForCausalLM  # Mengimpor tokenizer dan model dari HuggingFace

# 1. Load tokenizer dan model dari HuggingFace Hub
model_name = "wahdie11/text2py2"           # Ganti dengan path ke model Anda di HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, ignore_mismatched_sizes=True)

# 2. Definisikan fungsi untuk menghasilkan kode dari input teks
def generate_code(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")  # Mengonversi prompt ke tensor
    outputs = model.generate(**inputs, max_length=256)  # Menghasilkan kode dengan panjang maksimum 256 token
    return tokenizer.decode(outputs[0], skip_special_tokens=True)  # Mendekode output menjadi string

# 3. Buat interface menggunakan Gradio
demo = gr.Interface(
    fn=generate_code,                        # Fungsi yang akan dijalankan
    inputs="text",                            # Input berupa teks
    outputs="text",                           # Output berupa teks
    title="Text-to-Code Generator",           # Judul aplikasi
    description="Masukkan deskripsi teks untuk menghasilkan kode."  # Deskripsi aplikasi
)

# 4. Menjalankan interface
demo.launch()