|
import streamlit as st |
|
import galai as gal |
|
|
|
|
|
|
|
@st.cache_resource |
|
def load_model(model_name): |
|
model = gal.load_model(model_name, num_gpus=0) |
|
return model |
|
|
|
if 'text' not in st.session_state: |
|
st.session_state['text'] = "" |
|
|
|
def generate_text(): |
|
|
|
st.session_state['text'] = model.generate(input_text=st.session_state.editor, penalty_alpha=penalty_alpha, top_k=top_k, max_new_tokens=max_new_tokens, new_doc=new_doc) |
|
|
|
def suggest_reference(): |
|
|
|
st.session_state['text'] = st.session_state.editor + " " + model.generate_reference(input_text=st.session_state.editor) |
|
|
|
|
|
st.sidebar.markdown("### Select model") |
|
|
|
choose_model = st.sidebar.selectbox("Size", ["mini", "base"]) |
|
|
|
model = load_model(choose_model) |
|
|
|
st.sidebar.markdown("### Text generation settings") |
|
|
|
max_new_tokens = st.sidebar.slider("Max new tokens", value=10, min_value = 10, max_value = 100, step=10) |
|
penalty_alpha = st.sidebar.slider("Alpha penalty", value = 0.6, min_value = 0.0, max_value=2.0, step=0.1) |
|
top_k = st.sidebar.slider("Top-k", min_value = 0, max_value=10, value = 2) |
|
new_doc = st.sidebar.checkbox("New document", value=True) |
|
|
|
|
|
st.markdown( |
|
''' |
|
# Scientific writing assistant |
|
|
|
## GALACTICA model |
|
|
|
The [GALACTICA models](https://www.galactica.org) have been training on a large corpus of scientific data (see also the [GitHub repository](https://github.com/paperswithcode/galai)). Try out the two smaller models here and how they can be used to generate scientific text and suggest references. |
|
|
|
Write text in the editor and push **TAB** or **CTRL+ENTER** to generate text. |
|
|
|
Settings for model size (mini = 125 M and base = 1.3 B parameters) and text generation can be managed using from the left margin. |
|
|
|
## Suggest citations |
|
|
|
Use **Add citation** button to suggest and insert a citation into the text editor. The citation format is *Title, First author*. |
|
''') |
|
|
|
add_ref = st.button("Add citation", on_click=suggest_reference) |
|
|
|
text_editor = st.text_area("(not shown)", st.session_state['text'], height = 500, key="editor", on_change=generate_text, label_visibility="collapsed") |
|
|
