Spaces:
Sleeping
Sleeping
import json | |
import os | |
import shutil | |
import requests | |
import gradio as gr | |
from huggingface_hub import Repository | |
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 = os.environ.get("API_URL") | |
with open("./HHH_prompt.txt", "r") as f: | |
HHH_PROMPT = f.read() + "\n\n" | |
FIM_PREFIX = "<fim_prefix>" | |
FIM_MIDDLE = "<fim_middle>" | |
FIM_SUFFIX = "<fim_suffix>" | |
FIM_INDICATOR = "<FILL_HERE>" | |
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: | |
``` | |
<reponame>REPONAME<filename>FILENAME<gh_stars>STARS\ncode<|endoftext|> | |
``` | |
STARS can be one of: 0, 1-10, 10-100, 100-1000, 1000+ | |
### Commits | |
The commits data is formatted as follows: | |
``` | |
<commit_before>code<commit_msg>text<commit_after>code<|endoftext|> | |
``` | |
### Jupyter structure | |
Jupyter notebooks were both trained in form of Python scripts as well as the following structured format: | |
``` | |
<start_jupyter><jupyter_text>text<jupyter_code>code<jupyter_output>output<jupyter_text> | |
``` | |
### Issues | |
We also trained on GitHub issues using the following formatting: | |
``` | |
<issue_start><issue_comment>text<issue_comment>...<issue_closed> | |
``` | |
### Fill-in-the-middle | |
Fill in the middle requires rearranging the model inputs. The playground does this for you - all you need is to specify where to fill: | |
``` | |
code before<FILL_HERE>code after | |
``` | |
""" | |
theme = gr.themes.Monochrome( | |
primary_hue="indigo", | |
secondary_hue="blue", | |
neutral_hue="slate", | |
radius_size=gr.themes.sizes.radius_sm, | |
font=[gr.themes.GoogleFont("Open Sans"), "ui-sans-serif", "system-ui", "sans-serif"], | |
) | |
client = Client( | |
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, chat_mode=False): | |
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, | |
seed=42, | |
) | |
if chat_mode: | |
generate_kwargs.update({"stop_sequences": ["\nHuman", "\n-----"]}) | |
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: | |
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 | |
previous_token = "" | |
for response in stream: | |
if fim_mode and response.token.text =="<|endoftext|>": | |
output += (suffix + "\n" + response.token.text) | |
elif chat_mode and response.token.text in ["Human", "-----"] and previous_token=="\n": | |
return output | |
else: | |
output += response.token.text | |
previous_token = response.token.text | |
yield output | |
return output | |
examples = [ | |
"def print_hello_world():", | |
'def fibonacci(n: int) -> int:\n """ Compute the n-th Fibonacci number. """', | |
"class TransformerDecoder(nn.Module):", | |
"class ComplexNumbers:" | |
] | |
def process_example(args): | |
for x in generate(args): | |
pass | |
return x | |
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 | |
_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.\ | |
""" | |
) | |
with gr.Row(): | |
with gr.Column(scale=3): | |
instruction = gr.Textbox(placeholder="Enter your prompt here", label="Prompt", elem_id="q-input") | |
submit = gr.Button("Generate", variant="primary") | |
output = gr.Code(elem_id="q-output") | |
with gr.Group(elem_id="share-btn-container"): | |
community_icon = gr.HTML(community_icon_html, visible=True) | |
loading_icon = gr.HTML(loading_icon_html, visible=True) | |
share_button = gr.Button("Share to community", elem_id="share-btn", visible=True) | |
gr.Examples( | |
examples=examples, | |
inputs=[instruction], | |
cache_examples=False, | |
fn=process_example, | |
outputs=[output], | |
) | |
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, | |
minimum=0.0, | |
maximum=2.0, | |
step=0.1, | |
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", | |
) | |
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", | |
) | |
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) |