| import streamlit as st |
| from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed |
| from transformers import pipeline |
| import torch |
| import json |
| import pandas as pd |
|
|
| @st.cache(allow_output_mutation=True) |
| def load_tokenizer(model_ckpt): |
| return AutoTokenizer.from_pretrained(model_ckpt) |
|
|
| @st.cache(allow_output_mutation=True) |
| def load_model(model_ckpt): |
| model = AutoModelForCausalLM.from_pretrained(model_ckpt, low_cpu_mem_usage=True) |
| return model |
|
|
| @st.cache() |
| def load_examples(): |
| with open("examples.json", "r") as f: |
| examples = json.load(f) |
| return examples |
|
|
| st.set_page_config(page_icon=':laptop:', layout="wide") |
| |
|
|
| st.sidebar.header("Models") |
| models = ["CodeParrot", "OPT", "InCoder"] |
| selected_models = st.sidebar.multiselect('Select code generation models to compare:', |
| models, |
| default=["CodeParrot"]) |
| st.sidebar.header("Tasks") |
| tasks = [" ","Model architecture", "Model evaluation", "Pretraining dataset", "Code generation"] |
| selected_task = st.sidebar.selectbox("Select a task:", tasks) |
|
|
|
|
| tokenizer1 = load_tokenizer("lvwerra/codeparrot") |
| model1 = load_model("lvwerra/codeparrot") |
| tokenizer2 = load_tokenizer("facebook/incoder-1B") |
| model2 = load_model("facebook/incoder-1B") |
| |
| |
| pipelines = {} |
| for element in models: |
| if element == "CodeParrot": |
| pipelines[element] = pipeline("text-generation", model=model1, tokenizer=tokenizer1) |
| elif element == "InCoder": |
| tokenizer = load_tokenizer("facebook/incoder-1B") |
| model = load_model("facebook/incoder-1B") |
| pipelines[element] = pipeline("text-generation", model=model2, tokenizer=tokenizer2) |
| |
| |
| |
| |
| |
| examples = load_examples() |
| example_names = [example["name"] for example in examples] |
| name2id = dict([(name, i) for i, name in enumerate(example_names)]) |
| set_seed(42) |
| gen_kwargs = {} |
|
|
| if selected_task == " ": |
| st.title("Code Generation Models comparison") |
| with open("intro.txt", "r") as f: |
| intro = f.read() |
| st.markdown(intro) |
| elif selected_task == "Pretraining dataset": |
| st.title("Pretraining datasets 📚") |
| st.markdown("Preview of some code files from Github repositories") |
| df = pd.read_csv("preview-github-data.csv") |
| st.dataframe(df) |
| for model in selected_models: |
| with open(f"datasets/{model.lower()}.txt", "r") as f: |
| text = f.read() |
| st.markdown(f"### {model}:") |
| st.markdown(text) |
| elif selected_task == "Model architecture": |
| st.title("Model architecture 🔨") |
| for model in selected_models: |
| with open(f"architectures/{model.lower()}.txt", "r") as f: |
| text = f.read() |
| st.markdown(f"## {model}:") |
| st.markdown(text) |
| elif selected_task == "Code generation": |
| st.title("Code generation 💻") |
| st.sidebar.header("Examples") |
| selected_example = st.sidebar.selectbox("Select one of the following examples:", example_names) |
| example_text = examples[name2id[selected_example]]["value"] |
| default_length = examples[name2id[selected_example]]["length"] |
| st.sidebar.header("Generation settings") |
| gen_kwargs["do_sample"] = st.sidebar.radio("Decoding strategy:", ["Greedy", "Sample"]) == "Sample" |
| gen_kwargs["max_new_tokens"] = st.sidebar.slider("Number of tokens to generate:", value=default_length, min_value=8, step=8, max_value=256) |
| if gen_kwargs["do_sample"]: |
| gen_kwargs["temperature"] = 0.2 |
| gen_kwargs["top_k"] = 0 |
| gen_kwargs["top_p"] = 0.95 |
| gen_prompt = st.text_area("Generate code with prompt:", value=example_text, height=220,).strip() |
| if st.button("Generate code!"): |
| with st.spinner("Generating code..."): |
| for model in selected_models: |
| if model != "OPT": |
| pipe = pipelines[model] |
| generated_text = pipe(gen_prompt, **gen_kwargs)[0]['generated_text'] |
| st.markdown(f"{model}:") |
| st.code(generated_text) |
|
|
