Add function type annotations.

app.py

@@ -6,9 +6,9 @@ import warnings
 import gradio as gr
 import requests
 import torch
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline, Pipeline
 
-from templates import starting_app_code, update_iframe_js, copy_snippet_js, download_code_js, load_js
+from templates import starting_app_code, update_iframe_js, copy_snippet_js, download_code_js, load_js, DemoType
 
 # Filter the UserWarning raised by the audio component.
 warnings.filterwarnings("ignore", message='Trying to convert audio automatically from int32 to 16-bit int format')
@@ -30,7 +30,7 @@ API_URL = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-b
 headers = {"Authorization": f"Bearer {HF_TOKEN}"}
 
 
-def init_speech_to_text_model():
+def init_speech_to_text_model() -> Pipeline:
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
     torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
@@ -56,19 +56,16 @@ whisper_pipe = init_speech_to_text_model()
 code_pattern = re.compile(r'```python\n(.*?)```', re.DOTALL)
 
 
-def query(payload):
+def query(payload: dict):
     response = requests.post(API_URL, headers=headers, json=payload)
     return response.json()
 
 
-def generate_text(code, prompt):
+def generate_text(code: str, prompt: str) -> (str, str, str):
     logger.info(f"Calling API with prompt:\n{prompt}")
     prompt = f"```python\n{code}```\nGiven the code above return only updated code for the following request:\n{prompt}\n<|assistant|>"
     params = {"max_new_tokens": 512}
-    output = query({
-        "inputs": prompt,
-        "parameters": params,
-    })
+    output = query({"inputs": prompt, "parameters": params})
     if 'error' in output:
         logger.warning(f'Language model call failed: {output["error"]}')
         raise gr.Warning(f'Language model call failed: {output["error"]}')
@@ -82,7 +79,7 @@ def generate_text(code, prompt):
     return assistant_reply, new_code, None
 
 
-def transcribe(audio):
+def transcribe(audio: str) -> (str, str):
     result = whisper_pipe(audio)
     return result["text"], None
 
@@ -100,21 +97,22 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 gr.Markdown("## 1. Run your app in the browser!")
-                html = gr.HTML(value='<div id="stliteDemoDiv"></div>')
+                gr.HTML(value='<div id="stliteDemoDiv"></div>')
         gr.Markdown("## 2. Customize using voice requests!")
         with gr.Row():
             with gr.Column():
                 with gr.Group():
                     in_audio = gr.Audio(label="Record a voice request", source='microphone', type='filepath')
                     in_prompt = gr.Textbox(label="Or type a text request and press Enter",
-                                           placeholder="Need an idea? Try one of these:\n- Add a button\n- Change the greeting from hello to hey there")
+                                           placeholder="Need an idea? Try one of these:\n- Add a button to reverse the name\n- Change the greeting to Spanish\n- Make the button primary")
                 out_text = gr.TextArea(label="Chat Assistant Response")
                 clear_btn = gr.ClearButton([in_prompt, in_audio, out_text])
             with gr.Column():
                 code_area = gr.Code(label="App Code - You can also edit directly and then click Update App",
-                                    language='python', value=starting_app_code('stlite'))
+                                    language='python', value=starting_app_code(DemoType.STREAMLIT))
                 update_btn = gr.Button("Update App", variant="primary")
-                code_update_params = {'fn': None, 'inputs': code_area, 'outputs': None, '_js': update_iframe_js('stlite')}
+                code_update_params = {'fn': None, 'inputs': code_area, 'outputs': None,
+                                      '_js': update_iframe_js(DemoType.STREAMLIT)}
                 gen_text_params = {'fn': generate_text, 'inputs': [code_area, in_prompt],
                                    'outputs': [out_text, code_area]}
                 transcribe_params = {'fn': transcribe, 'inputs': [in_audio], 'outputs': [in_prompt, in_audio]}
@@ -125,9 +123,9 @@ with gr.Blocks() as demo:
             with gr.Column():
                 gr.Markdown("## 3. Export your app to share!")
                 copy_snippet_btn = gr.Button("Copy app snippet to paste in another page")
-                copy_snippet_btn.click(copy_notify, code_area, None, _js=copy_snippet_js('stlite'))
+                copy_snippet_btn.click(copy_notify, code_area, None, _js=copy_snippet_js(DemoType.STREAMLIT))
                 download_btn = gr.Button("Download app as a standalone file")
-                download_btn.click(None, code_area, None, _js=download_code_js('stlite'))
+                download_btn.click(None, code_area, None, _js=download_code_js(DemoType.STREAMLIT))
         with gr.Row():
             with gr.Column():
                 gr.Markdown("## Current limitations")
@@ -138,7 +136,7 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 gr.Markdown("## 1. Run your app in the browser!")
-                html = gr.HTML(value='<div id="gradioDemoDiv"></div>')
+                gr.HTML(value='<div id="gradioDemoDiv"></div>')
         gr.Markdown("## 2. Customize using voice requests!")
         with gr.Row():
             with gr.Column():
@@ -150,11 +148,12 @@ with gr.Blocks() as demo:
                 clear = gr.ClearButton([in_prompt, in_audio, out_text])
             with gr.Column():
                 code_area = gr.Code(label="App Code - You can also edit directly and then click Update App",
-                                    language='python', value=starting_app_code('gradio-lite'))
+                                    language='python', value=starting_app_code(DemoType.GRADIO))
                 update_btn = gr.Button("Update App", variant="primary")
                 code_update_params = {'fn': None, 'inputs': code_area, 'outputs': None,
-                                      '_js': update_iframe_js('gradio-lite')}
-                gen_text_params = {'fn': generate_text, 'inputs': [code_area, in_prompt], 'outputs': [out_text, code_area]}
+                                      '_js': update_iframe_js(DemoType.GRADIO)}
+                gen_text_params = {'fn': generate_text, 'inputs': [code_area, in_prompt],
+                                   'outputs': [out_text, code_area]}
                 transcribe_params = {'fn': transcribe, 'inputs': [in_audio], 'outputs': [in_prompt, in_audio]}
                 update_btn.click(**code_update_params)
                 in_prompt.submit(**gen_text_params).then(**code_update_params)
@@ -163,9 +162,9 @@ with gr.Blocks() as demo:
             with gr.Column():
                 gr.Markdown("## 3. Export your app to share!")
                 copy_snippet_btn = gr.Button("Copy app snippet to paste in another page")
-                copy_snippet_btn.click(copy_notify, code_area, None, _js=copy_snippet_js('gradio-lite'))
+                copy_snippet_btn.click(copy_notify, code_area, None, _js=copy_snippet_js(DemoType.GRADIO))
                 download_btn = gr.Button("Download app as a standalone file")
-                download_btn.click(None, code_area, None, _js=download_code_js('gradio-lite'))
+                download_btn.click(None, code_area, None, _js=download_code_js(DemoType.GRADIO))
         with gr.Row():
             with gr.Column():
                 gr.Markdown("## Current limitations")
@@ -173,9 +172,10 @@ with gr.Blocks() as demo:
                     gr.Markdown(
                         "- Only gradio-lite apps using the python standard libraries and gradio are supported\n- The chat hasn't been tuned on gradio library data; it may make mistakes\n- The app needs to fully reload each time it is changed")
 
-    stlite_tab.select(lambda: "stlite", None, selectedTab).then(None, None, None, _js=load_js('stlite'))
-    gradio_lite_tab.select(lambda: "gradio-lite", None, selectedTab).then(None, None, None, _js=load_js('gradio-lite'))
-    demo.load(None, None, None, _js=load_js('stlite'))
+    stlite_tab.select(lambda: "stlite", None, selectedTab).then(None, None, None, _js=load_js(DemoType.STREAMLIT))
+    gradio_lite_tab.select(lambda: "gradio-lite", None, selectedTab).then(None, None, None,
+                                                                          _js=load_js(DemoType.GRADIO))
+    demo.load(None, None, None, _js=load_js(DemoType.STREAMLIT))
     demo.css = "footer {visibility: hidden}"
 
 if __name__ == "__main__":

templates.py

@@ -1,20 +1,27 @@
+from enum import Enum
 from pathlib import Path
 
+
+class DemoType(Enum):
+    GRADIO = 1
+    STREAMLIT = 2
+
+
 gradio_lite_html_template = Path('templates/gradio-lite/gradio-lite-template.html').read_text()
 stlite_html_template = Path('templates/stlite/stlite-template.html').read_text()
 stlite_snippet_template = Path('templates/stlite/stlite-snippet-template.html').read_text()
 
 
-def starting_app_code(demo_type):
-    if demo_type == 'gradio-lite':
+def starting_app_code(demo_type: DemoType) -> str:
+    if demo_type == DemoType.GRADIO:
         return Path('templates/gradio-lite/gradio_lite_starting_code.py').read_text()
-    elif demo_type == 'stlite':
+    elif demo_type == DemoType.STREAMLIT:
         return Path('templates/stlite/stlite_starting_code.py').read_text().replace('`', r'\`')
     raise NotImplementedError(f'{demo_type} is not a supported demo type')
 
 
-def load_js(demo_type):
-    if demo_type == 'gradio-lite':
+def load_js(demo_type: DemoType) -> str:
+    if demo_type == DemoType.GRADIO:
         return f"""() => {{
             if (window.gradioLiteLoaded) {{
                 return
@@ -32,7 +39,7 @@ def load_js(demo_type):
             frame.contentWindow.document.close();
             window.gradioLiteLoaded = true;
         }}"""
-    elif demo_type == 'stlite':
+    elif demo_type == DemoType.STREAMLIT:
         return f"""() => {{
             if (window.stliteLoaded) {{
                 return
@@ -54,8 +61,8 @@ def load_js(demo_type):
     raise NotImplementedError(f'{demo_type} is not a supported demo type')
 
 
-def update_iframe_js(demo_type):
-    if demo_type == 'gradio-lite':
+def update_iframe_js(demo_type: DemoType) -> str:
+    if demo_type == DemoType.GRADIO:
         #  TODO: Works but is inefficient because the iframe has to be reloaded each time
         return f"""(code) => {{
             const pattern = /# APP CODE START(.*?)# APP CODE END/gs;
@@ -77,7 +84,7 @@ def update_iframe_js(demo_type):
             frame.contentWindow.document.write(completedTemplate);
             frame.contentWindow.document.close();
         }}"""
-    elif demo_type == 'stlite':
+    elif demo_type == DemoType.STREAMLIT:
         return f"""async (code) => {{       
             async function update() {{
                 const appController = document.getElementById('demo-iframe').contentWindow.window.appController;
@@ -90,8 +97,8 @@ def update_iframe_js(demo_type):
     raise NotImplementedError(f'{demo_type} is not a supported demo type')
 
 
-def copy_snippet_js(demo_type):
-    if demo_type == 'gradio-lite':
+def copy_snippet_js(demo_type: DemoType) -> str:
+    if demo_type == DemoType.GRADIO:
         return f"""async (code) => {{
             const pattern = /# APP CODE START(.*?)# APP CODE END/gs;
             const template = `<div id="KiteWindApp">\n<script type="module" crossorigin src="https://cdn.jsdelivr.net/npm/@gradio/lite/dist/lite.js"></script>
@@ -103,7 +110,7 @@ def copy_snippet_js(demo_type):
             const snippet = completedTemplate;
             await navigator.clipboard.writeText(snippet);
         }}"""
-    elif demo_type == 'stlite':
+    elif demo_type == DemoType.STREAMLIT:
         return f"""async (code) => {{
             const escapedCode = code.replace('`', String.fromCharCode(92) + '`');
             const template = `{stlite_html_template}`;
@@ -116,8 +123,8 @@ def copy_snippet_js(demo_type):
     raise NotImplementedError(f'{demo_type} is not a supported demo type')
 
 
-def download_code_js(demo_type):
-    if demo_type == 'gradio-lite':
+def download_code_js(demo_type: DemoType) -> str:
+    if demo_type == demo_type.GRADIO:
         return f"""(code) => {{
             const pattern = /# APP CODE START(.*?)# APP CODE END/gs;
             const template = `{gradio_lite_html_template}`;
@@ -141,7 +148,7 @@ def download_code_js(demo_type):
             // Clean up by revoking the URL
             URL.revokeObjectURL(url);
         }}"""
-    elif demo_type == 'stlite':
+    elif demo_type == demo_type.STREAMLIT:
         return f"""(code) => {{
             const escapedCode = code.replace('`', String.fromCharCode(92) + '`');
             // Step 1: Generate the HTML content