Update app.py

app.py

@@ -12,6 +12,12 @@ HF_TOKEN = os.environ.get("HF_TOKEN", None)
 API_URL = os.environ.get("API_URL")
 
 
+FIM_PREFIX = "<fim_prefix>"
+FIM_MIDDLE = "<fim_middle>"
+FIM_SUFFIX = "<fim_suffix>"
+
+FIM_INDICATOR = "<FILL_HERE>"
+
 theme = gr.themes.Monochrome(
     primary_hue="indigo",
     secondary_hue="blue",
@@ -25,34 +31,44 @@ client = Client(
     #headers={"Authorization": f"Bearer {HF_TOKEN}"},
 )
 
-
 def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
 
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
     top_p = float(top_p)
-
+    fim_mode = False
+    
     generate_kwargs = dict(
         temperature=temperature,
         max_new_tokens=max_new_tokens,
         top_p=top_p,
         repetition_penalty=repetition_penalty,
         do_sample=True,
-        truncate=999,
         seed=42,
-        stop_sequences=["</s>"],
     )
 
-    stream = client.generate_stream(
-        prompt,
-        **generate_kwargs,
-    )
+    if FIM_INDICATOR in prompt:
+        fim_mode = True
+        try:
+            prefix, suffix = prompt.split("<FILL-HERE>")
+        except:
+            ValueError("Only one <FILL-HERE> allowed in prompt!")
+        prompt = f"{FIM_PREFIX}{prefix}{FIM_SUFFIX}{suffix}{FIM_MIDDLE}"
+
+    stream = client.generate_stream(prompt, **generate_kwargs)
+
+    if fim_mode:
+        output = prefix
+    else:
+        output = prompt
 
-    output = prompt
     for response in stream:
         output += response.token.text
         yield output
+    
+    if fim_mode:
+        output += suffix
     return output
 
 
@@ -74,7 +90,33 @@ css = ".generating {visibility: hidden}" + share_btn_css
 with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
     with gr.Column():
         gr.Markdown(
-            """ # BigCode - Playground
+            """\
+            # BigCode - Playground
+            _Note:_ this is an internal playground - please do not share. The deployment can also change and thus the space not work as we continue development.
+
+            ## Model formats
+
+            ### Prefixes
+            Any combination of the three:
+            ```
+            <reponame>REPONAME<filename>FILENAME<gh_stars>STARS\nCode<eos>
+            ```
+            Stars be: 0, 1-10, 10-100, 100-1000, 1000+
+
+            ### Commits
+            ```
+            <commit_before>code<commit_msg>text<commit_after>code<|endoftext|>
+            ```
+            
+            ### Jupyter structure
+            ```
+            <start_jupyter><jupyter_text>text<jupyter_code>code<jupyter_output>output<jupyter_text>
+            ```
+
+            ### Fill-in-the-middle
+            ```
+            code before<FILL_HERE>code after
+            ```
             """
         )
         with gr.Row():
@@ -107,7 +149,7 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
                     label="Max new tokens",
                     value=256,
                     minimum=0,
-                    maximum=512,
+                    maximum=4096,
                     step=4,
                     interactive=True,
                     info="The maximum numbers of new tokens",