reshinthadith/code-representation-learning

1import streamlit as st2from transformers import AutoModelForCausalLM, AutoTokenizer3 4 5MODEL_NAME = "reshinthadith/BashGPTNeo"6def load_model_and_tokenizer(model_name):7    """Adding load_model_and_tokenizer function to keep the model in the memory"""8    model = AutoModelForCausalLM.from_pretrained(model_name)9    tokenizer = AutoTokenizer.from_pretrained(model_name)10    return tokenizer,model11 12tokenizer,model = load_model_and_tokenizer(MODEL_NAME)13 14MAX_TOKS = 12815MAX_NEW_TOKS = 128 16def generate_text(prompt):17    prompt = "<english> " + prompt + " <bash>"18    inputs = tokenizer(prompt, truncation=True, return_tensors="pt")19    output_seq = model.generate(20        input_ids=inputs.input_ids, max_length=MAX_TOKS,21        max_new_tokens=MAX_NEW_TOKS,22        do_sample=True, temperature=0.8,23        num_return_sequences=124    )25 26    outputs = tokenizer.batch_decode(output_seq, skip_special_tokens=True)27    outputs = outputs[0].split("<bash>")[-1]28    return outputs29st.set_page_config(30page_title= "Code Representation Learning",31    32    initial_sidebar_state= "expanded"33    )34 35st.sidebar.title("Code Representation Learning")36st.sidebar.write("work by Reshinth Adithyan & Aditya Thuruvas")37 38workflow = st.sidebar.selectbox('select a task', ['Bash Synthesis'])39if workflow == "Bash Synthesis":40    st.title("Program Synthesis for Bash")41    prompt = st.text_input("Natural Language prompt ",'print all the files with ".cpp" extension')42    button = st.button("synthesize")43if button:44    generated_text = generate_text(prompt)45    st.write(generated_text)46