Update app.py

app.py

@@ -1,31 +1,103 @@
-# Step 1: Improved layout and clean sidebar
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+Enhanced Gradio UI for the Salesforce/codet5-large model using the Hugging Face Inference API.
+Adheres to best practices, PEP8, flake8, and the Zen of Python.
+"""
+
 import gradio as gr
 
 MODEL_ID = "Salesforce/codet5-large"
 
-with gr.Blocks(fill_height=True, theme=gr.themes.Soft()) as demo:
-    with gr.Sidebar():
-        gr.Markdown("## 🤖 Inference Provider")
-        gr.Markdown(
-            f"This Space showcases the `{MODEL_ID}` model, served via the Hugging Face Inference API.\n\n"
-            "🔐 Sign in with your Hugging Face account to use this API."
-        )
-        login_button = gr.LoginButton("🔐 Sign in")
 
-        gr.Markdown("---")
-        gr.Markdown(f"**Model:** `{MODEL_ID}`")
-        gr.Markdown("[📄 View Model Card](https://huggingface.co/Salesforce/codet5-large)")
+def prepare_payload(prompt: str, max_tokens: int) -> dict:
+    """
+    Prepare the payload dictionary for the Hugging Face inference call.
+
+    Args:
+        prompt (str): The input code containing `<extra_id_0>`.
+        max_tokens (int): Maximum number of tokens for generation.
+
+    Returns:
+        dict: Payload for the model API call.
+    """
+    return {"inputs": prompt, "parameters": {"max_length": max_tokens}}
+
+
+def extract_generated_text(api_response: dict) -> str:
+    """
+    Extract generated text from the API response.
+
+    Args:
+        api_response (dict): The response dictionary from the model API call.
+
+    Returns:
+        str: The generated text, or string representation of the response.
+    """
+    return api_response.get("generated_text", str(api_response))
 
-    gr.Markdown("# 🧠 CodeT5 Inference UI")
 
-    # Placeholder: we'll add inputs, buttons, outputs in next steps
-    gr.Markdown("_Interface under construction..._")
+def main():
+    with gr.Blocks(fill_height=True, theme=gr.themes.Soft()) as demo:
+        with gr.Sidebar():
+            gr.Markdown("## 🤖 Inference Provider")
+            gr.Markdown(
+                (
+                    "This Space showcases the `{}` model, served via the Hugging Face Inference API.\n\n"
+                    "Sign in with your Hugging Face account to access the model."
+                ).format(MODEL_ID)
+            )
+            login_button = gr.LoginButton("🔐 Sign in")
+            gr.Markdown("---")
+            gr.Markdown(f"**Model:** `{MODEL_ID}`")
+            gr.Markdown("[📄 View Model Card](https://huggingface.co/Salesforce/codet5-large)")
+
+        gr.Markdown("# 🧠 CodeT5 Inference UI")
+        gr.Markdown("Enter your Python code snippet with `<extra_id_0>` as the mask token.")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                code_input = gr.Code(
+                    label="Input Code",
+                    language="python",
+                    value="def greet(user): print(f'hello <extra_id_0>!')",
+                    lines=10,
+                    autofocus=True,
+                )
+                max_tokens = gr.Slider(
+                    minimum=8, maximum=128, value=32, step=8, label="Max Tokens"
+                )
+                submit_btn = gr.Button("🚀 Run Inference")
+            with gr.Column(scale=1):
+                output_text = gr.Textbox(
+                    label="Inference Output",
+                    lines=10,
+                    interactive=False,
+                    placeholder="Model output will appear here...",
+                )
+
+        # Load the model from Hugging Face Inference API.
+        model_iface = gr.load(
+            f"models/{MODEL_ID}",
+            accept_token=login_button,
+            provider="hf-inference",
+        )
+
+        # Chain click events: prepare payload -> API call -> extract output.
+        submit_btn.click(
+            fn=prepare_payload,
+            inputs=[code_input, max_tokens],
+            outputs=model_iface,
+            api_name="prepare_payload",
+        ).then(
+            fn=extract_generated_text,
+            inputs=model_iface,
+            outputs=output_text,
+            api_name="extract_output",
+        )
+
+    demo.launch()
 
-    # Load HF inference API model
-    gr.load(
-        f"models/{MODEL_ID}",
-        accept_token=login_button,
-        provider="hf-inference"
-    )
 
-demo.launch()
+if __name__ == "__main__":
+    main()