codeparrot/code-generation-models

1import json2import os3import pandas as pd4import requests5import threading6import streamlit as st7from datasets import load_dataset, load_metric8 9MODELS = ["CodeParrot", "InCoder", "CodeGen", "PolyCoder"]10GENERATION_MODELS = ["CodeParrot", "InCoder", "CodeGen"]11 12 13@st.cache()14def load_examples():15    with open("utils/examples.json", "r") as f:16        examples = json.load(f)17    return examples18    19    20def load_evaluation():21    # load task 2 of HumanEval and code_eval_metric22    os.environ["HF_ALLOW_CODE_EVAL"] = "1"23    human_eval = load_dataset("openai_humaneval")24    entry_point = f"check({human_eval['test'][2]['entry_point']})"25    test_func = "\n" + human_eval["test"][2]["test"] + "\n" + entry_point26    code_eval = load_metric("code_eval")27    return code_eval, test_func28 29 30def read_markdown(path):31    with open(path, "r") as f:32        output = f.read()33    st.markdown(output, unsafe_allow_html=True)34 35 36def generate_code(37    generations, model_name, gen_prompt, max_new_tokens, temperature, seed38):39    # call space using its API endpoint40    url = (41        f"https://hf.space/embed/codeparrot/{model_name.lower()}-subspace/+/api/predict/"42    )43    r = requests.post(44        url=url, json={"data": [gen_prompt, max_new_tokens, temperature, seed]}45    )46    generated_text = r.json()["data"][0]47    generations.append({model_name: generated_text})48 49 50def generate_code_threads(51    generations, models, gen_prompt, max_new_tokens, temperature, seed52):53    threads = []54    for model_name in models:55        # create the thread56        threads.append(57            threading.Thread(58                target=generate_code,59                args=(60                    generations,61                    model_name,62                    gen_prompt,63                    max_new_tokens,64                    temperature,65                    seed,66                ),67            )68        )69        threads[-1].start()70 71    for t in threads:72        t.join()73 74@st.cache(show_spinner=False)75def generate_teaser(gen_prompt):76    generations = []77    generate_code(generations, "CodeParrot", gen_prompt, 8, 0.2, 42)78    return generations[0]["CodeParrot"]79    80st.set_page_config(page_icon=":laptop:", layout="wide")81with open("utils/table_contents.md", "r") as f:82    contents = f.read()83 84st.sidebar.markdown(contents)85 86# Introduction87st.title("Code generation with 🤗")88read_markdown("utils/summary.md")89## teaser90example_text = "def print_hello_world():"91col1, col2, col3 = st.columns([1, 2, 1])92with col2:93    gen_prompt = st.text_area(94        "",95        value=example_text,96        height=100,97    ).strip()98    if st.button("Generate code!", key=1):99        with st.spinner("Generating code..."):100            st.code(generate_teaser(gen_prompt)) 101read_markdown("utils/intro.md")102 103# Code datasets104st.subheader("1 - Code datasets")105read_markdown("datasets/intro.md")106read_markdown("datasets/github_code.md")107col1, col2 = st.columns([1, 2])108with col1:109    selected_model = st.selectbox("", MODELS, key=1)110read_markdown(f"datasets/{selected_model.lower()}.md")111 112 113# Model architecture114st.subheader("2 - Model architecture")115read_markdown("architectures/intro.md")116col1, col2 = st.columns([1, 2])117with col1:118    selected_model = st.selectbox("", MODELS, key=2)119read_markdown(f"architectures/{selected_model.lower()}.md")120 121# Model evaluation122st.subheader("3 - Code model evaluation")123read_markdown("evaluation/intro.md")124read_markdown("evaluation/demo_humaneval.md")125## quiz126st.markdown("Below you can try solving this problem or visualize the solution of CodeParrot:")127with open("evaluation/problem.md", "r") as f:128    problem = f.read()129with open("evaluation/solution.md", "r") as f:130    solution = f.read()131    132candidate_solution = st.text_area(133    "Complete the problem:",134    value=problem,135    height=240,136).strip()137if st.button("Test my solution", key=2):138    with st.spinner("Testing..."):139        code_eval, test_func = load_evaluation()140        test_cases = [test_func]141        candidates = [[candidate_solution]]142        pass_at_k, _ = code_eval.compute(references=test_cases, predictions=candidates)143        text = "Your solution didn't pass the test, pass@1 is 0 😕" if pass_at_k['pass@1'] < 1  else "Congrats your pass@1 is 1! 🎉"144        st.markdown(text)145if st.button("Show model solution", key=3):146    st.markdown(solution)147    148# Code generation149st.subheader("4 - Code generation ✨")150read_markdown("generation/intro.md")151col1, col2, col3 = st.columns([7, 1, 6])152with col1:153    st.markdown("**Models**")154    selected_models = st.multiselect(155        "Select code generation models to compare:",156        GENERATION_MODELS,157        default=GENERATION_MODELS,158        key=3,159    )160    st.markdown(" ")161    st.markdown("**Examples**")162    examples = load_examples()163    example_names = [example["name"] for example in examples]164    name2id = dict([(name, i) for i, name in enumerate(example_names)])165    selected_example = st.selectbox(166        "Select one of the following examples or implement yours:", example_names167    )168    example_text = examples[name2id[selected_example]]["value"]169    default_length = examples[name2id[selected_example]]["length"]170with col3:171    st.markdown("**Generation settings**")172    temperature = st.slider(173        "Temperature:", value=0.2, min_value=0.1, step=0.1, max_value=2.0174    )175    max_new_tokens = st.slider(176        "Number of tokens to generate:",177        value=default_length,178        min_value=8,179        step=4,180        max_value=256,181    )182    seed = st.slider("Random seed:", value=42, min_value=0, step=1, max_value=1000)183gen_prompt = st.text_area(184    "Generate code with prompt:",185    value=example_text,186    height=200,187).strip()188if st.button("Generate code!", key=4):189    with st.spinner("Generating code..."):190        # use threading191        generations = []192        generate_code_threads(193            generations,194            selected_models,195            gen_prompt=gen_prompt,196            max_new_tokens=max_new_tokens,197            temperature=temperature,198            seed=seed,199        )200        for i in range(len(generations)):201            st.markdown(f"**{selected_models[i]}**")202            for j in range(len(generations)):203                if selected_models[i] in generations[j].keys():204                    st.code(generations[j][selected_models[i]])205        if len(generations) < len(selected_models):206            st.markdown("<span style='color:red'>Warning: Some models run into timeout, try another time or reduce the Number of tokens to generate. You can also try generating code using the original subspaces: [InCoder](https://huggingface.co/spaces/loubnabnl/incoder-subspace), [CodeGen](https://huggingface.co/spaces/loubnabnl/codegen-subspace), [CodeParrot](https://huggingface.co/spaces/loubnabnl/codeparrot-subspace)</span>", unsafe_allow_html=True)207 208# Resources209st.subheader("Resources")210read_markdown("utils/resources.md")211