Update app.py

app.py

@@ -1,21 +1,31 @@
-import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import gradio as gr
 
-model_id = "OSS-Forge/codet5p-770m-vhdl"
+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)
+
+def generate_output(input_text):
+    inputs = tokenizer.encode(input_text, return_tensors='pt')
+    outputs = model.generate(
+        inputs,
+        max_length=256,
+        num_beams=5,
+        early_stopping=True,
+    )
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
+
 
 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",
+    fn=generate_output,
+    inputs=gr.Textbox(lines=5, placeholder='Insert the English description here...'),
+    outputs=gr.Textbox(),
+    title='VHDL Code Generator (CodeT5+ 770M)',
+    description='Generate VHDL code from an English description using a fine-tuned CodeT5+ model.'
 )
 
-iface.launch()
+
+iface.launch()