from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

# CodeT5 modelini yükle
model_name = "Salesforce/codet5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Kod üretme fonksiyonu
def generate_java_code(prompt):
    # Girdi metnini modele uygun bir şekilde biçimlendir
    prompt = f"Generate a complete Java program:\n{prompt}"
    
    # Girdiyi tokenize et ve modele gönder
    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
    outputs = model.generate(
        inputs["input_ids"],
        max_length=300,  # Çıktının uzunluğunu artır
        num_beams=5,     # Beam search ile kaliteyi artır
        early_stopping=True
    )
    # Tokenize edilmiş çıktıyı kod olarak döndür
    code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Post-processing: Eksik yapıları düzelt
    if "public static void main" not in code:
        code = f"public class Main {{\n{code}\n}}"
    if code.strip().endswith("{"):
        code += "\n    // Complete the logic here\n}"
    
    return code

# Gradio UI tasarımı
with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align: center;'>Java Kod Üretici (CodeT5)</h1>")
    prompt = gr.Textbox(
        label="Doğal Dil Girdisi",
        placeholder="Örnek: Write a Java program to find the larger of two numbers."
    )
    output_code = gr.Textbox(
        label="Üretilen Java Kodu",
        lines=15,  # Çıktı alanı için daha fazla satır
    )
    btn = gr.Button("Kod Üret")
    btn.click(generate_java_code, inputs=prompt, outputs=output_code)

demo.launch()