|
import platform |
|
from typing import Tuple, List, Dict, Optional |
|
|
|
import streamlit as st |
|
import torch |
|
from trecover.config import var |
|
from trecover.utils.beam_search import beam_search, dashboard_loop |
|
from trecover.utils.inference import data_to_columns, create_noisy_columns |
|
from trecover.utils.transform import columns_to_tensor, tensor_to_target |
|
from trecover.utils.visualization import visualize_columns, visualize_target |
|
|
|
MAX_CHARS = 256 |
|
|
|
PLAIN_EXAMPLES = { |
|
'Select example': None, |
|
'Example 1': 'As people around the country went into the streets to cheer the conviction, some businesses in ' |
|
'Portland boarded up their windows once again.', |
|
'Example 2': 'That night, a small group of activists wearing black approached a group of journalists, threatening' |
|
' to smash the cameras of those who remained on scene.', |
|
'Example 3': 'English as we know it today came to be exported to other parts of the world through British ' |
|
'colonisation, and is now the dominant language in Britain' |
|
} |
|
|
|
NOISED_EXAMPLES = { |
|
'Select example': None, |
|
'Example 1': 'a ds fpziq ofe ngkhbo p pghl ue waq frlqjo o u dnxrm dgr yrtsco kho deuasm dhysc ao u nwzhy tle r ' |
|
'yzpe xwabc gce nger klqto wiq nfprso t no tpgq tcfh ae twas tw ur re e t gyutsm t xgo rc ubhq e wle ' |
|
'r ty h nwpeaq xdsc o dnhelm v thir ikcq tkuo i o twn ps frio mo oe b kuiqtb jsq zi tnye ge dgrqs s ' |
|
'cioe ys whic wne wp thlo dnprsc xvpyrt hurlm kveaj nbfp dome pbeaj dusmo a r dzrqsm xace du nxkuai ' |
|
'gpulcm tpi h pie uim r wbhrj ui n dwgp dkeio nkwhqs zs' |
|
} |
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
model = torch.hub.load('alex-snd/TRecover', model='trecover', device=device, version='latest') |
|
|
|
|
|
def main() -> None: |
|
st.set_page_config( |
|
page_title='TRecover', |
|
page_icon='🩹', |
|
layout='wide', |
|
initial_sidebar_state='expanded') |
|
|
|
if 'history' not in st.session_state: |
|
st.session_state.history = list() |
|
|
|
if 'data' not in st.session_state: |
|
st.session_state.data = '' |
|
|
|
if 'regenerate' not in st.session_state: |
|
st.session_state.regenerate = False |
|
|
|
if 'columns' not in st.session_state: |
|
st.session_state.columns = None |
|
|
|
if 'is_unix' not in st.session_state: |
|
st.session_state.is_unix = platform.system() != 'Windows' |
|
|
|
sidebar() |
|
|
|
|
|
def set_regenerate() -> None: |
|
st.session_state.regenerate = True |
|
|
|
|
|
def unset_regenerate() -> None: |
|
st.session_state.regenerate = False |
|
|
|
|
|
def sidebar() -> None: |
|
st.sidebar.markdown(body= |
|
""" |
|
<h1 align="center"> |
|
<font size="20">🤷</font> |
|
<a href="https://alex-snd.github.io/TRecover">About the Project</a> |
|
</h1> |
|
<br><br> |
|
""", |
|
unsafe_allow_html=True) |
|
|
|
option = st.sidebar.radio('Sections', ('Inference', 'Inference history')) |
|
|
|
if option == 'Inference': |
|
is_plain, min_noise, max_noise, bw = inference_sidebar() |
|
inference_page(is_plain, min_noise, max_noise, bw) |
|
else: |
|
history_sidebar() |
|
history_page() |
|
|
|
|
|
def inference_sidebar() -> Tuple[bool, int, int, int]: |
|
st.sidebar.text('\n') |
|
|
|
data_type = st.sidebar.radio('Input type', ('Plain text', 'Noisy columns'), key='data_type', |
|
index=0 if 'Plain text' == st.session_state.get('data_type', 'Plain text') else 1) |
|
is_plain = data_type == 'Plain text' |
|
|
|
st.sidebar.text('\n') |
|
|
|
if is_plain: |
|
min_noise, max_noise = st.sidebar.slider('\nNoise range', 0, 5, key='noise_range', |
|
value=st.session_state.get('noise_range', (0, 5)), |
|
on_change=set_regenerate) |
|
else: |
|
min_noise, max_noise = 0, 0 |
|
|
|
bw = st.sidebar.slider('Beam search width', 1, 25, key='beam_width', |
|
value=st.session_state.get('beam_width', 5)) |
|
|
|
if max_noise > var.MAX_NOISE: |
|
st.sidebar.warning('Max noise value is too large. This will entail poor performance') |
|
|
|
return is_plain, min_noise, max_noise + 1, bw |
|
|
|
|
|
def history_sidebar() -> None: |
|
pass |
|
|
|
|
|
def save_to_history(is_plain: bool, |
|
min_noise: int, |
|
max_noise: int, |
|
bw: int, |
|
columns: List[str], |
|
chains: List[Tuple[str, float]] |
|
) -> None: |
|
text = st.session_state.data if is_plain else None |
|
|
|
st.session_state.history.append((is_plain, text, min_noise, max_noise, bw, columns, chains)) |
|
|
|
|
|
@st.cache(ttl=3600, show_spinner=False, suppress_st_warning=True) |
|
def predict(columns: List[str], bw: int) -> List[Tuple[str, float]]: |
|
src = columns_to_tensor(columns, device) |
|
|
|
chains = beam_search(src, model, bw, device, beam_loop=dashboard_loop) |
|
chains = [(visualize_target(tensor_to_target(chain)), prob) for (chain, prob) in chains] |
|
|
|
return chains |
|
|
|
|
|
def get_noisy_columns(data: str, min_noise: int, max_noise: int) -> List[str]: |
|
columns = create_noisy_columns(data, min_noise, max_noise) |
|
|
|
return [''.join(set(c)) for c in columns] |
|
|
|
|
|
def get_input_data(examples: Dict[str, Optional[str]], max_chars: int) -> str: |
|
input_field, examples_filed = st.columns([1, 0.27]) |
|
|
|
option = examples_filed.selectbox(label='', options=examples.keys()) |
|
|
|
return input_field.text_input(label='', value=examples[option] or st.session_state.data, max_chars=max_chars) |
|
|
|
|
|
def inference_page(is_plain: bool, min_noise: int, max_noise: int, bw: int) -> None: |
|
st.subheader('Insert plain text' if is_plain else 'Insert noisy columns separated by spaces') |
|
|
|
if is_plain: |
|
data = get_input_data(PLAIN_EXAMPLES, max_chars=MAX_CHARS) |
|
else: |
|
data = get_input_data(NOISED_EXAMPLES, max_chars=MAX_CHARS * 4) |
|
|
|
if not data: |
|
st.stop() |
|
|
|
if is_plain: |
|
if st.session_state.regenerate or not st.session_state.columns or data != st.session_state.data: |
|
columns = get_noisy_columns(data, min_noise, max_noise) |
|
st.session_state.columns = columns |
|
unset_regenerate() |
|
else: |
|
columns = st.session_state.columns |
|
else: |
|
columns = data_to_columns(data, separator=' ') |
|
|
|
st.session_state.data = data |
|
|
|
st.subheader('\nColumns') |
|
st.text(visualize_columns(columns, delimiter='')) |
|
st.subheader('\n') |
|
|
|
placeholder = st.empty() |
|
recover_field, regen_filed = placeholder.columns([.11, 1]) |
|
|
|
if is_plain: |
|
regen_filed.button('Regenerate', on_click=set_regenerate) |
|
|
|
if columns and recover_field.button('Recover'): |
|
if st.session_state.is_unix: |
|
with placeholder.container(): |
|
progress_bar_placeholder = st.empty() |
|
st.button('Stop') |
|
|
|
with progress_bar_placeholder: |
|
chains = predict(columns, bw) |
|
else: |
|
with placeholder: |
|
chains = predict(columns, bw) |
|
|
|
with placeholder.container(): |
|
st.subheader('\nPrediction') |
|
st.text('\n\n'.join(chain for chain, _ in chains)) |
|
|
|
if st.button('Clear'): |
|
st.session_state.task_id = None |
|
|
|
save_to_history(is_plain, min_noise, max_noise, bw, columns, chains) |
|
|
|
|
|
def history_page() -> None: |
|
st.header('Inference History') |
|
|
|
if len(st.session_state.history) == 0: |
|
st.info('No records saved') |
|
return |
|
|
|
for record_id, (is_plain, text, min_noise, max_noise, bw, columns, chains) in enumerate(st.session_state.history, |
|
start=1): |
|
st.info(f'Record {record_id}') |
|
|
|
if is_plain: |
|
st.text(f'Plain data: {text}') |
|
|
|
st.text(f'Noise range: [{min_noise}, {max_noise}]') |
|
st.text(f'Beam search width: {bw}') |
|
|
|
st.text('Columns:') |
|
st.text(visualize_columns(columns, delimiter='')) |
|
|
|
st.text('Prediction:') |
|
st.text('\n\n'.join(chain for chain, _ in chains)) |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
