trl-lib/stack-llama

1import json2import os3import shutil4 5import gradio as gr6from huggingface_hub import Repository, CommitScheduler7from text_generation import Client8 9from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css10 11HF_TOKEN = os.environ.get("TRL_TOKEN", None)12API_URL = "https://api-inference.huggingface.co/models/kashif/stack-llama-2"13 14 15theme = gr.themes.Monochrome(16    primary_hue="indigo",17    secondary_hue="blue",18    neutral_hue="slate",19    radius_size=gr.themes.sizes.radius_sm,20    font=[gr.themes.GoogleFont("Open Sans"), "ui-sans-serif", "system-ui", "sans-serif"],21)22if HF_TOKEN:23    try:24        shutil.rmtree("./data/")25    except:26        pass27 28    # Schedule regular uploads every 10 minutes. Remote repo and local folder are created if they don't already exist.29    scheduler = CommitScheduler(30        repo_id="trl-lib/stack-llama-2-prompts",31        repo_type="dataset",32        folder_path="./data/",33        path_in_repo="./",34        every=10,35        token=HF_TOKEN36    )37 38 39client = Client(40    API_URL,41    headers={"Authorization": f"Bearer {HF_TOKEN}"},42)43 44PROMPT_TEMPLATE = """Question: {prompt}\n\nAnswer:"""45 46 47def save_inputs_and_outputs(inputs, outputs, generate_kwargs):48    with open(os.path.join("data", "prompts.jsonl"), "a") as f:49        json.dump({"inputs": inputs, "outputs": outputs, "generate_kwargs": generate_kwargs}, f, ensure_ascii=False)50        f.write("\n")51 52def generate(instruction, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0,  do_save=True):53    formatted_instruction = PROMPT_TEMPLATE.format(prompt=instruction)54 55    temperature = float(temperature)56    if temperature < 1e-2:57        temperature = 1e-258    top_p = float(top_p)59 60    generate_kwargs = dict(61        temperature=temperature,62        max_new_tokens=max_new_tokens,63        top_p=top_p,64        repetition_penalty=repetition_penalty,65        do_sample=True,66        truncate=999,67        seed=42,68        stop_sequences=["</s>"],69    )70 71    stream = client.generate_stream(72        formatted_instruction,73        **generate_kwargs,74    )75 76    output = ""77    for response in stream:78        output += response.token.text79        yield output80    if HF_TOKEN and do_save:81        try:82            print("Pushing prompt and completion to the Hub")83            save_inputs_and_outputs(formatted_instruction, output, generate_kwargs)84        except Exception as e:85            print(e)86            87    return output88 89 90examples = [91    "A llama is in my lawn. How do I get rid of him?",92    "What are the various algorithms to sort a list?",93    "How can I sort a list in Python?",94    "How do I ask a question in StackOverflow?",95    "How to beat a Hitmonlee in a Pokemon battle?",96    "How can I write a Java function to generate the nth Fibonacci number?",97]98 99 100def process_example(args):101    for x in generate(args):102        pass103    return x104 105css = ".generating {visibility: hidden}" + share_btn_css106 107with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:108    with gr.Column():109        gr.Markdown(110            """![](https://huggingface.co/spaces/trl-lib/stack-llama/resolve/main/stackllama_logo.png)111 112 113            StackLLaMa-2 is a 7 billion parameter language model based on [Meta's LLaMA 2 model](https://ai.meta.com/llama/) that has been trained on pairs of questions and answers from [Stack Exchange](https://stackexchange.com) using Direct Preference Optimization (DPO) with the [TRL library](https://github.com/lvwerra/trl). For more details, check out our [blog post](https://huggingface.co/blog/dpo-trl).114 115            Type in the box below and click the button to generate answers to your most pressing questions!116 117            ⚠️ **Intended Use**: this app and its [supporting model](https://huggingface.co/kashif/stack-llama-2) are provided as educational tools to explain RLHF with the TRL library; not to serve as replacement for human expertise. For more details on the model's limitations in terms of factuality and biases, see the [model card.](https://huggingface.co/kashif/stack-llama-2#intended-uses--limitations)118            119            ⚠️ **Data Collection**: by default, we are collecting the prompts entered in this app to further improve and evaluate the model. Do not share any personal or sensitive information while using the app! You can opt out of this data collection by removing the checkbox below:120      """121        )122        with gr.Row():123            with gr.Column(scale=3):124                do_save = gr.Checkbox(125                        value=True,126                        label="Store data",127                        info="You agree to the storage of your prompt and generated text for research and development purposes:")128                instruction = gr.Textbox(placeholder="Enter your question here", label="Question", elem_id="q-input")129 130                131                with gr.Box():132                    gr.Markdown("**Answer**")133                    output = gr.Markdown(elem_id="q-output")134                submit = gr.Button("Generate", variant="primary")135                with gr.Group(elem_id="share-btn-container"):136                    community_icon = gr.HTML(community_icon_html, visible=True)137                    loading_icon = gr.HTML(loading_icon_html, visible=True)138                    share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)139                gr.Examples(140                    examples=examples,141                    inputs=[instruction],142                    cache_examples=False,143                    fn=process_example,144                    outputs=[output],145                )146 147            with gr.Column(scale=1):148                149                temperature = gr.Slider(150                    label="Temperature",151                    value=0.9,152                    minimum=0.0,153                    maximum=2.0,154                    step=0.1,155                    interactive=True,156                    info="Higher values produce more diverse outputs",157                )158                max_new_tokens = gr.Slider(159                    label="Max new tokens",160                    value=256,161                    minimum=0,162                    maximum=512,163                    step=4,164                    interactive=True,165                    info="The maximum numbers of new tokens",166                )167                top_p = gr.Slider(168                    label="Top-p (nucleus sampling)",169                    value=0.90,170                    minimum=0.0,171                    maximum=1,172                    step=0.05,173                    interactive=True,174                    info="Higher values sample more low-probability tokens",175                )176                repetition_penalty = gr.Slider(177                    label="Repetition penalty",178                    value=1.2,179                    minimum=1.0,180                    maximum=2.0,181                    step=0.05,182                    interactive=True,183                    info="Penalize repeated tokens",184                )185 186    submit.click(generate, inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty, do_save], outputs=[output])187    instruction.submit(generate, inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty], outputs=[output])188    share_button.click(None, [], [], _js=share_js)189 190demo.queue(concurrency_count=16).launch(debug=True)