Structured dataclasses for requests and settings.

Grouping client and utils source files into their own src directory

README.md

@@ -1,5 +1,5 @@
 ---
-title: StarCoder Completion
+title: StarCoder Demo
 emoji: 💫
 colorFrom: gray
 colorTo: blue
@@ -11,15 +11,17 @@ duplicated_from: bigcode/bigcode-playground
 ---
 
 
-# StarCoder Demo
+# ⭐StarCoder Demo💫
 
-## Star💫Coder Code Completion Playground 💻
+## Code-Completion Playground 💻 with ⭐StarCoder Models
 
-This is a demo playground to generate code with the power of [StarCoder](https://huggingface.co/bigcode/starcoder)💫 a **15B** parameter model for code generation in **80+** programming languages.
+This is a demo playground to generate code with the power of ⭐[StarCoder](https://huggingface.co/bigcode/starcoder) a **15B** parameter model for code generation in **80+** programming languages.
 
-This is not an instruction model but just a code completion tool.
+ℹ️ This is not an instruction model but just a code completion tool.
 
-For instruction and chatting you can chat with a prompted version of the model directly at the [HuggingFace Chat 🤗 (hf.co/chat)](https://huggingface.co/chat/?model=starcoder)
+🗣️For instruction and chatting you can chat with a prompted version of the model directly at the [HuggingFace🤗Chat💬(hf.co/chat)](https://huggingface.co/chat/?model=starcoder)
+
+![StarCoder](https://huggingface.co/datasets/bigcode/admin/resolve/main/StarCoderBanner.png)
 
 ---
 

app.py

@@ -1,27 +1,33 @@
 import sys
 import os
+import logging as log
+from typing import Generator
 
 import gradio as gr
 from gradio.themes.utils import sizes
 from text_generation import Client
+from src.request import StarCoderRequest, StarCoderRequestConfig
 
 # todo: remove and replace by the actual js file instead
-from share_btn import (share_js)
-from utils import (
+from src.share_btn import (share_js)
+from src.utils import (
     get_file_as_string,
     get_sections,
     get_url_from_env_or_default_path,
     preview
 )
 from constants import (
-    DEFAULT_STARCODER_API_PATH,
-    DEFAULT_STARCODER_BASE_API_PATH,
     FIM_MIDDLE,
     FIM_PREFIX,
     FIM_SUFFIX,
     END_OF_TEXT,
     MIN_TEMPERATURE,
 )
+from settings import (
+    DEFAULT_PORT,
+    DEFAULT_STARCODER_API_PATH,
+    DEFAULT_STARCODER_BASE_API_PATH,
+)
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 # Gracefully exit the app if the HF_TOKEN is not set,
@@ -37,17 +43,13 @@ if not HF_TOKEN:
     # gr.close_all(verbose=False)
     sys.exit(1)
 
-API_URL = get_url_from_env_or_default_path("STARCODER_API", DEFAULT_STARCODER_API_PATH)
+API_URL_STAR = get_url_from_env_or_default_path("STARCODER_API", DEFAULT_STARCODER_API_PATH)
 API_URL_BASE = get_url_from_env_or_default_path("STARCODER_BASE_API", DEFAULT_STARCODER_BASE_API_PATH)
 
-preview("StarCoder Model's URL", API_URL)
-preview("StarCoderBase Model's URL", API_URL_BASE)
+preview("StarCoder Model URL", API_URL_STAR)
+preview("StarCoderBase Model URL", API_URL_BASE)
 preview("HF Token", HF_TOKEN, ofuscate=True)
 
-DEFAULT_PORT = 7860
-
-FIM_INDICATOR = "<FILL_HERE>"
-
 # Loads the whole content of the formats.md file
 # and stores it into the FORMATS variable
 STATIC_PATH = "static"
@@ -81,55 +83,74 @@ theme = gr.themes.Monochrome(
 HEADERS = {
     "Authorization": f"Bearer {HF_TOKEN}",
 }
-client = Client(API_URL, headers = HEADERS)
-client_base = Client(API_URL_BASE, headers = HEADERS)
+client_star = Client(API_URL_STAR, headers=HEADERS)
+client_base = Client(API_URL_BASE, headers=HEADERS)
 
-def generate(prompt,
+def get_tokens_collector(request: StarCoderRequest) -> Generator[str, None, None]:
+
+    model_client = client_star if request.settings.version == "StarCoder" else client_base
+    stream = model_client.generate_stream(request.prompt, **request.settings.kwargs())
+    for response in stream:
+        # print(response.token.id, response.token.text)
+        # if token.text != END_OF_TEXT:
+        if response.token.id != 0:
+            yield response.token.text
+
+def get_tokens_accumulator(request: StarCoderRequest) -> Generator[str, None, None]:
+    # start with the prefix (if in fim_mode)
+    output = request.prefix if request.fim_mode else request.prompt
+    for token in get_tokens_collector(request=request):
+        output += token
+        yield output
+    # after the last token, append the suffix (if in fim_mode)
+    if request.fim_mode:
+        output += request.suffix
+        yield output
+    # Append an extra line at the end
+    yield output + '\n'
+
+def get_tokens_linker(request: StarCoderRequest) -> str:
+    return "".join(list(get_tokens_collector(request)))
+
+def generate(
+        prompt: str,
         temperature = 0.9,
         max_new_tokens = 256,
         top_p = 0.95,
         repetition_penalty = 1.0,
         version = "StarCoder",
-    ):
-
-    temperature = min(float(temperature), MIN_TEMPERATURE)
-    top_p = float(top_p)
-
-    generate_kwargs = dict(
-        temperature = temperature,
-        max_new_tokens = max_new_tokens,
-        top_p = top_p,
-        repetition_penalty = repetition_penalty,
-        do_sample = True,
-        seed = 42,
+    ) -> Generator[str, None, None]:
+    request = StarCoderRequest(
+        prompt=prompt,
+        settings=StarCoderRequestConfig(
+            version=version,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+        )
     )
+    yield from get_tokens_accumulator(request)
 
-    if fim_mode := FIM_INDICATOR in prompt:
-        try:
-            prefix, suffix = prompt.split(FIM_INDICATOR)
-        except Exception as err:
-            print(str(err))
-            raise ValueError(f"Only one {FIM_INDICATOR} allowed in prompt!") from err
-        prompt = f"{FIM_PREFIX}{prefix}{FIM_SUFFIX}{suffix}{FIM_MIDDLE}"
-
-    model_client = client if version == "StarCoder" else client_base
-
-    stream = model_client.generate_stream(prompt, **generate_kwargs)
-
-    output = prefix if fim_mode else prompt
-
-    for response in stream:
-        if response.token.text == END_OF_TEXT:
-            if fim_mode:
-                output += suffix
-            else:
-                return output
-        else:
-            output += response.token.text
-        # todo: log this value while in debug mode
-        # previous_token = response.token.text
-        yield output
-    return output
+def process_example(
+        prompt: str,
+        temperature = 0.9,
+        max_new_tokens = 256,
+        top_p = 0.95,
+        repetition_penalty = 1.0,
+        version = "StarCoder",
+    ) -> Generator[str, None, None]:
+    request = StarCoderRequest(
+        prompt=prompt,
+        settings=StarCoderRequestConfig(
+            version=version,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+        )
+    )
+    yield from get_tokens_linker(request)
 
 # todo: move it into the README too
 examples = [
@@ -138,12 +159,6 @@ examples = [
     "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",
 ]
 
-def process_example(args):
-    for x in generate(args):
-        pass
-    return x
-
-
 with gr.Blocks(theme=theme, analytics_enabled=False, css=CSS) as demo:
     with gr.Column():
         gr.Markdown(description)
@@ -226,6 +241,9 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=CSS) as demo:
         generate,
         inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty, version],
         outputs=[output],
+        # preprocess=False,
+        max_batch_size=8,
+        show_progress=True
     )
     share_button.click(None, [], [], _js=share_js)
 

constants.py

@@ -1,10 +1,7 @@
-DEFAULT_HUGGINGFACE_MODELS_API_BASE_URL = "https://api-inference.huggingface.co/models/"
-DEFAULT_STARCODER_API_PATH = "bigcode/starcoder/"
-DEFAULT_STARCODER_BASE_API_PATH = "bigcode/starcoderbase/"
-
 FIM_PREFIX = "<fim_prefix>"
 FIM_MIDDLE = "<fim_middle>"
 FIM_SUFFIX = "<fim_suffix>"
 END_OF_TEXT = "<|endoftext|>"
 
-MIN_TEMPERATURE = 1e-3
+# Near zero temperature to avoid division by zero
+MIN_TEMPERATURE = 1e-4

settings.py

@@ -0,0 +1,14 @@
+# URLs for the StarCoder Models/APIs
+DEFAULT_HUGGINGFACE_MODELS_API_BASE_URL = "https://api-inference.huggingface.co/models/"
+DEFAULT_STARCODER_API_PATH = "bigcode/starcoder/"
+DEFAULT_STARCODER_BASE_API_PATH = "bigcode/starcoderbase/"
+FIM_INDICATOR = "<FILL_HERE>"
+DEFAULT_PORT = 7860
+
+DEFAULT_SETTINGS = dict(
+    temperature = 0.9,
+    max_new_tokens = 256,
+    top_p = 0.95,
+    repetition_penalty = 1.0,
+    version = "StarCoder",
+)

src/request.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+from typing import Dict, Any, Union
+
+from constants import (
+    FIM_MIDDLE,
+    FIM_PREFIX,
+    FIM_SUFFIX,
+    MIN_TEMPERATURE,
+)
+from settings import (
+    FIM_INDICATOR,
+)
+
+@dataclass
+class StarCoderRequestConfig:
+    temperature: float
+    max_new_tokens: int
+    top_p: float
+    repetition_penalty: float
+    version: str
+
+    def __post_init__(self):
+        self.temperature = min(float(self.temperature), MIN_TEMPERATURE)
+        self.max_new_tokens = int(self.max_new_tokens)
+        self.top_p = float(self.top_p)
+        self.repetition_penalty = float(self.repetition_penalty)
+        self.do_sample = True
+        self.seed = 42
+
+    def __repr__(self) -> str:
+        """Returns a custom string representation of the Configurations."""
+        values = dict(
+            model = self.version,
+            temp = self.temperature,
+            tokens = self.max_new_tokens,
+            p = self.top_p,
+            penalty = self.repetition_penalty,
+            sample = self.do_sample,
+            seed = self.seed,
+        )
+        return f"StarCoderRequestConfig({values})"
+
+    def kwargs(self) -> Dict[str, Union[Any, float, int]]:
+        """
+        Returns a custom dictionary representation of the Configurations.
+        removing the model version.
+        """
+        values = vars(self).copy()
+        values.pop("version")
+        return values
+
+@dataclass
+class StarCoderRequest:
+    prompt: str
+    settings: StarCoderRequestConfig
+
+    def __post_init__(self):
+        self.fim_mode = FIM_INDICATOR in self.prompt
+        self.prefix, self.suffix = None, None
+        if self.fim_mode:
+            try:
+                self.prefix, self.suffix = self.prompt.split(FIM_INDICATOR)
+            except Exception as err:
+                print(str(err))
+                raise ValueError(f"Only one {FIM_INDICATOR} allowed in prompt!") from err
+            self.prompt = f"{FIM_PREFIX}{self.prefix}{FIM_SUFFIX}{self.suffix}{FIM_MIDDLE}"
+
+    def __repr__(self) -> str:
+        """Returns a custom string representation of the Request."""
+        values = dict(
+            prompt = self.prompt,
+            configuration = self.settings,
+        )
+        return f"StarCoderRequest({values})"

share_btn.py → src/share_btn.py

@@ -67,7 +67,7 @@ ${outputTxt}`;
         .map(([key, value]) => `${encodeURIComponent(key)}=${encodeURIComponent(value)}`)
         .join('&');
 
-	window.open(`https://huggingface.co/spaces/Fisharp/starcoder-playground/discussions/new?${paramsStr}`, '_blank');
+	window.open(`https://huggingface.co/spaces/fisharp/starcoder-playground/discussions/new?${paramsStr}`, '_blank');
 
     shareBtnEl.style.removeProperty('pointer-events');
     shareIconEl.style.removeProperty('display');

utils.py → src/utils.py

@@ -2,7 +2,7 @@ import os
 from typing import List
 from urllib.parse import urljoin
 
-from constants import DEFAULT_HUGGINGFACE_MODELS_API_BASE_URL
+from settings import DEFAULT_HUGGINGFACE_MODELS_API_BASE_URL
 
 
 def masked(value: str, n_shown: int, length: int = None) -> str:

static/share_btn.js

@@ -67,7 +67,7 @@ ${outputTxt}`;
         .map(([key, value]) => `${encodeURIComponent(key)}=${encodeURIComponent(value)}`)
         .join('&');
 
-    window.open(`https://huggingface.co/spaces/Fisharp/starcoder-playground/discussions/new?${paramsStr}`, '_blank');
+    window.open(`https://huggingface.co/spaces/fisharp/starcoder-playground/discussions/new?${paramsStr}`, '_blank');
 
     shareBtnEl.style.removeProperty('pointer-events');
     shareIconEl.style.removeProperty('display');