add api endpoints and dropdown for models

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
@@ -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)