flax-community/code-clippy-problem-solver

1import urllib2 3import streamlit as st4from transformers import AutoModelForCausalLM, AutoTokenizer5 6# model_name = "flax-community/gpt-neo-1.3B-apps-all"7model_name = "flax-community/gpt-neo-125M-apps-all"8 9 10@st.cache(allow_output_mutation=True, max_entries=1)11def get_model():12    model = AutoModelForCausalLM.from_pretrained(model_name)13    tokenizer = AutoTokenizer.from_pretrained(model_name)14    tokenizer.pad_token = tokenizer.eos_token15    return (model, tokenizer)16 17 18def format_input(question, starter_code=""):19    answer_type = (20        "\nUse Call-Based format\n" if starter_code else "\nUse Standard Input format\n"21    )22    return f"\nQUESTION:\n{question}\n{starter_code}\n{answer_type}\nANSWER:\n"23 24 25def clean_text(generation):26    # clean up text has discussed in OpenAI's paper "Evaluating Large Language Models Trained on Code"27    generation = generation.split("\ndef")[0]28    generation = generation.split("\nclass")[0]29    generation = generation.split("\n#")[0]30    generation = generation.split("\nif")[0]31 32    return generation33 34 35def generate_solution(36    model, tokenizer, question, starter_code="", temperature=1.0, num_beams=137):38    prompt = format_input(question, starter_code)39    input_ids = tokenizer(prompt, return_tensors="pt").input_ids40    start = len(input_ids[0])41 42    output = model.generate(43        input_ids,44        max_length=start + 150,45        do_sample=True,46        top_p=0.95,47        pad_token_id=tokenizer.pad_token_id,48        eos_token_id=tokenizer.eos_token_id,49        early_stopping=True,50        temperature=temperature,51        num_beams=int(num_beams),52        no_repeat_ngram_size=None,53        repetition_penalty=None,54        num_return_sequences=None,55    )56    output_str = tokenizer.decode(output[0][start:], skip_special_tokens=True).strip()57    output_str = clean_text(output_str)58 59    return output_str60 61 62_EXAMPLES = [63    [64        """65Given a 2D list of size `m * n`. Your task is to find the sum of minimum value in each row.66For Example:67```python68[69  [1, 2, 3, 4, 5],       # minimum value of row is 170  [5, 6, 7, 8, 9],       # minimum value of row is 571  [20, 21, 34, 56, 100]  # minimum value of row is 2072]73```74So, the function should return `26` because sum of minimums is as `1 + 5 + 20 = 26`75        """,76        "",77        0.8,78    ],79    [80        """81# Personalized greeting82 83Create a function that gives a personalized greeting. This function takes two parameters: `name` and `owner`.84        """,85        """86Use conditionals to return the proper message:87 88case| return89--- | ---90name equals owner | 'Hello boss'91otherwise         | 'Hello guest'92def greet(name, owner):93        """,94        0.8,95    ],96]97 98 99def run():100    st.set_page_config(page_title="Code Clippy Problem Solver")101    # sidebar102    st.sidebar.title("Code Clippy")103    st.sidebar.image(104        "https://raw.githubusercontent.com/ncoop57/gpt-code-clippy/camera-ready/code_clippy_logo.jpg",105        caption="(c) awesome Aimee Trevett",106    )107    st.sidebar.markdown("[Github](https://github.com/ncoop57/gpt-code-clippy)")108    st.sidebar.markdown("[Report](https://github.com/ncoop57/gpt-code-clippy/wiki)")109 110    st.sidebar.markdown("### Controls:")111 112    temperature = st.sidebar.slider(113        "Temperature",114        min_value=0.5,115        max_value=1.5,116        value=0.8,117        step=0.1,118    )119    num_beams = st.sidebar.slider(120        "Num beams",121        min_value=1,122        max_value=4,123        step=1,124    )125 126    # main body127    model, tokenizer = get_model()128 129    question = st.text_input(130        "Problem: ",131        value="A function that can greet user by name. Given a name it should say hello to user.",132        help="Text description of the coding problem to be solved",133    )134    starter_code = st.text_input(135        "Started code: ", value="def greet(name):", help="Optional starter code"136    )137    submit_button = st.button("Solve")138 139    if submit_button:140        text = st.text("Generating solution...")141        # gif from https://giphy.com/gifs/alan-DfSXiR60W9MVq142        gif_runner = st.image("./loading.gif")143        output = generate_solution(144            model, tokenizer, question, starter_code, temperature, num_beams145        )146        text.empty()147        gif_runner.empty()148 149        st.text("Solution:")150        st.code(output, language="python")151 152        # Create link to carbon to make a nice screenshot of the generated code153        url_code = urllib.parse.quote(f"# {question}\n{output}")154        st.markdown(155            f"[Would you like a Carbon Copy?](https://carbon.now.sh/?bg=rgba%280%2C0%2C0%2C0%29&t=seti&wt=none&l=python&ds=false&dsyoff=20px&dsblur=68px&wc=true&wa=false&pv=56px&ph=56px&ln=false&fl=1&fm=Hack&fs=14px&lh=133%25&si=false&es=2x&wm=false&code={url_code})"156        )157 158 159if __name__ == "__main__":160    run()