Update app.py

app.py

@@ -11,6 +11,8 @@ from share_btn import community_icon_html, loading_icon_html, share_js, share_bt
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 API_URL = os.environ.get("API_URL")
 
+with open("https://huggingface.co/spaces/bigcode/bigcode-playground/raw/main/HHH_prompt.txt", "r") as f:
+    HHH_PROMPT = f.read() + "\n\n"
 
 FIM_PREFIX = "<fim_prefix>"
 FIM_MIDDLE = "<fim_middle>"
@@ -22,6 +24,9 @@ FORMATS = """## Model formats
 
 The model is pretrained on code and in addition to the pure code data it is formatted with special tokens. E.g. prefixes specifying the source of the file or special tokens separating code from a commit message. See below: 
 
+### Chat mode
+Chat mode prepends the [HHH prompt](https://gist.github.com/jareddk/2509330f8ef3d787fc5aaac67aab5f11#file-hhh_prompt-txt) from Anthropic to the request which conditions the model to be an assistant.
+
 ### Prefixes
 Any combination of the three following prefixes can be found in pure code files:
 
@@ -64,11 +69,10 @@ theme = gr.themes.Monochrome(
 )
 
 client = Client(
-    API_URL,
-    #headers={"Authorization": f"Bearer {HF_TOKEN}"},
+    API_URL,  #headers={"Authorization": f"Bearer {HF_TOKEN}"},
 )
 
-def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
+def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, chat_mode=False):
 
     temperature = float(temperature)
     if temperature < 1e-2:
@@ -85,18 +89,27 @@ def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition
         seed=42,
     )
 
+    if chat_mode and FIM_INDICATOR in prompt:
+        raise ValueError("Chat mode and FIM are mutually exclusive. Choose one or the other.")
+
+    if chat_mode:
+        chat_prompt = "Human: " + prompt + "\n\nAssistant:"
+        prompt = HHH_PROMPT + chat_prompt
+    
     if FIM_INDICATOR in prompt:
         fim_mode = True
         try:
             prefix, suffix = prompt.split(FIM_INDICATOR)
         except:
-            ValueError(f"Only one {FIM_INDICATOR} allowed in prompt!")
+            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 fim_mode:
         output = prefix
+    elif: chat_mode:
+        output = chat_prompt
     else:
         output = prompt
 
@@ -155,7 +168,11 @@ _Note:_ this is an internal playground - please do not share. The deployment can
                 gr.Markdown(FORMATS)
 
             with gr.Column(scale=1):
-                
+                chat_mode = gr.Checkbox(
+                    value=False,
+                    label="Chat mode",
+                    info="Uses Anthropic's HHH prompt to turn the model into an assistant."
+                )
                 temperature = gr.Slider(
                     label="Temperature",
                     value=0.2,
@@ -193,7 +210,7 @@ _Note:_ this is an internal playground - please do not share. The deployment can
                     info="Penalize repeated tokens",
                 )
 
-    submit.click(generate, inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty], outputs=[output])
-    instruction.submit(generate, inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty], outputs=[output])
+    submit.click(generate, inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty, chat_mode], outputs=[output])
+    # instruction.submit(generate, inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty, chat_mode], outputs=[output])
     share_button.click(None, [], [], _js=share_js)
 demo.queue(concurrency_count=16).launch(debug=True)