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") #tokenizer3 = load_tokenizer("facebook/opt-1.3b") #model3 = load_model("facebook/opt-1.3b") 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) #else: # tokenizer = load_tokenizer("facebook/opt-1.3b") # model = load_model("facebook/opt-1.3b") # pipelines[element] = pipeline("text-generation", model=model3, tokenizer=tokenizer3) 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)