Update app.py

app.py

@@ -10,8 +10,10 @@ from text_generation import Client
 from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css
 
 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"
+API_URL_PLUS = "https://api-inference.huggingface.co/models/bigcode/starcoderplus"
 
 FIM_PREFIX = "<fim_prefix>"
 FIM_MIDDLE = "<fim_middle>"
@@ -80,6 +82,9 @@ client = Client(
 client_base = Client(
     API_URL_BASE, headers={"Authorization": f"Bearer {HF_TOKEN}"},
 )
+client_plus = Client(
+    API_URL_PLUS, headers={"Authorization": f"Bearer {HF_TOKEN}"},
+)
 
 def generate(
     prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, version="StarCoder",
@@ -110,6 +115,8 @@ def generate(
 
     if version == "StarCoder":
         stream = client.generate_stream(prompt, **generate_kwargs)
+    elif version == "StarCoderPlus":
+        stream = client_plus.generate_stream(prompt, **generate_kwargs)
     else:
         stream = client_base.generate_stream(prompt, **generate_kwargs)
 
@@ -135,6 +142,7 @@ def generate(
 examples = [
     "X_train, y_train, X_test, y_test = train_test_split(X, y, test_size=0.1)\n\n# Train a logistic regression model, predict the labels on the test set and compute the accuracy score",
     "// Returns every other value in the array as a new array.\nfunction everyOther(arr) {",
+    "Poor English: She no went to the market. Corrected English:",
     "def alternating(list1, list2):\n   results = []\n   for i in range(min(len(list1), len(list2))):\n       results.append(list1[i])\n       results.append(list2[i])\n   if len(list1) > len(list2):\n       <FILL_HERE>\n   else:\n       results.extend(list2[i+1:])\n   return results",
 ]
 
@@ -159,9 +167,16 @@ css += share_btn_css + monospace_css + ".gradio-container {color: black}"
 
 description = """
 <div style="text-align: center;">
-    <h1> 💫 StarCoder<span style='color: #e6b800;'> - </span>Code Completion Playground 🪐</h1>
-    <p>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.
-    <b>This is not an instruction model</b>. For instruction and chatting you can chat with a fine-tuned version of the model at <a href="https://huggingface.co/spaces/HuggingFaceH4/starchat-playground">StarChat Playground</a></p>
+    <h1> ⭐ StarCoder <span style='color: #e6b800;'>Models</span> Playground</h1>
+    <p>This is a demo to generate text and code with the following StarCoder models:</p>
+</div>
+<div style="text-align: left;">
+    <ul>
+        <li><a href="https://huggingface.co/bigcode/starcoderplus" style='color: #e6b800;'>StarCoderPlus</a>: A finetuned version of StarCoderBase on English web data, making it strong in both English text and code generation.</li>
+        <li><a href="https://huggingface.co/bigcode/starcoderbase" style='color: #e6b800;'>StarCoderBase</a>: A code generation model trained on 80+ programming languages, providing broad language coverage for code generation tasks.</li>
+        <li><a href="https://huggingface.co/bigcode/starcoder" style='color: #e6b800;'>StarCoder</a>: A finetuned version of StarCoderBase specifically focused on Python, while also maintaining strong performance on other programming languages.</li>
+    </ul>
+    <p><b>Please note:</b> These models are not designed for instruction purposes. If you're looking for instruction or want to chat with a fine-tuned model, you can visit the <a href="https://huggingface.co/spaces/HuggingFaceH4/starchat-playground">StarChat Playground</a>.</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>\
@@ -170,15 +185,23 @@ disclaimer = """⚠️<b>Any use or sharing of this demo constitues your accepta
 with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
     with gr.Column():
         gr.Markdown(description)
+        with gr.Row():
+            version = gr.Dropdown(
+                        ["StarCoderPlus", "StarCoderBase", "StarCoder"],
+                        value="StarCoderPlus",
+                        label="Model",
+                        info="Choose a model from the list",
+                        )
         with gr.Row():
             with gr.Column():
                 instruction = gr.Textbox(
                     placeholder="Enter your code here",
-                    label="Code",
+                    lines=5,
+                    label="Input",
                     elem_id="q-input",
                 )
                 submit = gr.Button("Generate", variant="primary")
-                output = gr.Code(elem_id="q-output", lines=30)
+                output = gr.Code(elem_id="q-output", lines=30, label="Output")
                 with gr.Row():
                     with gr.Column():
                         with gr.Accordion("Advanced settings", open=False):
@@ -222,13 +245,7 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
                                         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)