add api endpoints and dropdown for models (#24)

- add api endpoints and dropdown for models (7af0e86df124c48b004a2ba74acd084b223b18fa)
- update readme to mention both models (ad689e69933185f3256c3e4e5cec1e4d5e6823a9)

app.py

@@ -11,7 +11,7 @@ from share_btn import community_icon_html, loading_icon_html, share_js, share_bt
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 API_URL = "https://api-inference.huggingface.co/models/bigcode/starcoder/"
-
+API_URL_BASE ="https://api-inference.huggingface.co/models/bigcode/starcoderbase/"
 
 FIM_PREFIX = "<fim_prefix>"
 FIM_MIDDLE = "<fim_middle>"
@@ -77,10 +77,12 @@ client = Client(
     API_URL,
     headers={"Authorization": f"Bearer {HF_TOKEN}"},
 )
-
+client_base = Client(
+    API_URL_BASE, headers={"Authorization": f"Bearer {HF_TOKEN}"},
+)
 
 def generate(
-    prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0
+    prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, version="StarCoder",
 ):
 
     temperature = float(temperature)
@@ -106,7 +108,10 @@ def generate(
             raise ValueError(f"Only one {FIM_INDICATOR} allowed in prompt!")
         prompt = f"{FIM_PREFIX}{prefix}{FIM_SUFFIX}{suffix}{FIM_MIDDLE}"
 
-    stream = client.generate_stream(prompt, **generate_kwargs)
+    if version == "StarCoder":
+        stream = client.generate_stream(prompt, **generate_kwargs)
+    else:
+        stream = client_base.generate_stream(prompt, **generate_kwargs)
 
     if fim_mode:
         output = prefix
@@ -160,7 +165,7 @@ css += share_btn_css + monospace_css + custom_output_css + ".gradio-container {c
 description = """
 <div style="text-align: center;">
     <h1 style='color: black;'> 💫 StarCoder<span style='color: #e6b800;'> - </span>Playground 🪐</h1>
-    <p style='color: black;'>This is a demo to generate code with <a href="https://huggingface.co/bigcode/starcoder" style='color: #e6b800;'>StarCoder</a>, a 15B parameter model for code generation in 86 programming languages.</p>
+    <p style='color: black;'>This is a demo to generate code with <a href="https://huggingface.co/bigcode/starcoder" style='color: #e6b800;'>StarCoder</a> and <a href="https://huggingface.co/bigcode/starcoderbase" style='color: #e6b800;'>StarCoderBase</a>, 15B parameter models for code generation in 86 programming languages.</p>
 </div>
 """
 disclaimer = """⚠️<b>Any use or sharing of this demo constitues your acceptance of the BigCode [OpenRAIL-M](https://huggingface.co/spaces/bigcode/bigcode-model-license-agreement) License Agreement and the use restrictions included within.</b>\
@@ -178,48 +183,56 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
                 )
                 submit = gr.Button("Generate", variant="primary")
                 output = gr.Code(elem_id="q-output", lines=30)
-
-                with gr.Accordion("Advanced settings", open=False):
-                    with gr.Row():
-                        column_1, column_2 = gr.Column(), gr.Column()
-                        with column_1:
-                            temperature = gr.Slider(
-                                label="Temperature",
-                                value=0.2,
-                                minimum=0.0,
-                                maximum=1.0,
-                                step=0.05,
-                                interactive=True,
-                                info="Higher values produce more diverse outputs",
-                            )
-                            max_new_tokens = gr.Slider(
-                                label="Max new tokens",
-                                value=256,
-                                minimum=0,
-                                maximum=8192,
-                                step=64,
-                                interactive=True,
-                                info="The maximum numbers of new tokens",
-                            )
-                        with column_2:
-                            top_p = gr.Slider(
-                                label="Top-p (nucleus sampling)",
-                                value=0.90,
-                                minimum=0.0,
-                                maximum=1,
-                                step=0.05,
-                                interactive=True,
-                                info="Higher values sample more low-probability tokens",
-                            )
-                            repetition_penalty = gr.Slider(
-                                label="Repetition penalty",
-                                value=1.2,
-                                minimum=1.0,
-                                maximum=2.0,
-                                step=0.05,
-                                interactive=True,
-                                info="Penalize repeated tokens",
-                            )
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Accordion("Advanced settings", open=False):
+                            with gr.Row():
+                                column_1, column_2 = gr.Column(), gr.Column()
+                                with column_1:
+                                    temperature = gr.Slider(
+                                        label="Temperature",
+                                        value=0.2,
+                                        minimum=0.0,
+                                        maximum=1.0,
+                                        step=0.05,
+                                        interactive=True,
+                                        info="Higher values produce more diverse outputs",
+                                    )
+                                    max_new_tokens = gr.Slider(
+                                        label="Max new tokens",
+                                        value=256,
+                                        minimum=0,
+                                        maximum=8192,
+                                        step=64,
+                                        interactive=True,
+                                        info="The maximum numbers of new tokens",
+                                    )
+                                with column_2:
+                                    top_p = gr.Slider(
+                                        label="Top-p (nucleus sampling)",
+                                        value=0.90,
+                                        minimum=0.0,
+                                        maximum=1,
+                                        step=0.05,
+                                        interactive=True,
+                                        info="Higher values sample more low-probability tokens",
+                                    )
+                                    repetition_penalty = gr.Slider(
+                                        label="Repetition penalty",
+                                        value=1.2,
+                                        minimum=1.0,
+                                        maximum=2.0,
+                                        step=0.05,
+                                        interactive=True,
+                                        info="Penalize repeated tokens",
+                                    )
+                    with gr.Column():
+                        version = gr.Dropdown(
+                                    ["StarCoderBase", "StarCoder"],
+                                    value="StarCoder",
+                                    label="Version",
+                                    info="",
+                                    )
                 gr.Markdown(disclaimer)
                 with gr.Group(elem_id="share-btn-container"):
                     community_icon = gr.HTML(community_icon_html, visible=True)
@@ -238,7 +251,7 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
 
     submit.click(
         generate,
-        inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty],
+        inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty, version],
         outputs=[output],
     )
     share_button.click(None, [], [], _js=share_js)