import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model_id = "OSS-Forge/codet5p-770m-vhdl"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id)

def generate_vhdl(description):
    inputs = tokenizer(description, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=256)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

iface = gr.Interface(
    fn=generate_vhdl,
    inputs=gr.Textbox(lines=4, label="Enter VHDL description"),
    outputs=gr.Textbox(lines=12, label="Generated VHDL"),
    title="VHDL Code Generator",
)

iface.launch()