import streamlit as st import torch import spacy # from spacy.lang.en import English # from utils import spacy_function, make_predictions, example_input from Dataset import SkimlitDataset from Embeddings import get_embeddings from Model import SkimlitModel from Tokenizer import Tokenizer from LabelEncoder import LabelEncoder from MakePredictions import make_skimlit_predictions from RandomAbstract import Choose_Random_text MODEL_PATH = 'skimlit-model-final-1.pt' TOKENIZER_PATH = 'tokenizer.json' LABEL_ENOCDER_PATH = "label_encoder.json" EMBEDDING_FILE_PATH = 'glove.6B.300d.txt' @st.cache() def create_utils(model_path, tokenizer_path, label_encoder_path, embedding_file_path): tokenizer = Tokenizer.load(fp=tokenizer_path) label_encoder = LabelEncoder.load(fp=label_encoder_path) embedding_matrix = get_embeddings(embedding_file_path, tokenizer, 300) model = SkimlitModel(embedding_dim=300, vocab_size=len(tokenizer), hidden_dim=128, n_layers=3, linear_output=128, num_classes=len(label_encoder), pretrained_embeddings=embedding_matrix) model.load_state_dict(torch.load(model_path, map_location='cpu')) print(model) return model, tokenizer, label_encoder def model_prediction(abstract, model, tokenizer, label_encoder): objective = '' background = '' method = '' conclusion = '' result = '' lines, pred = make_skimlit_predictions(abstract, model, tokenizer, label_encoder) # pred, lines = make_predictions(abstract) for i, line in enumerate(lines): if pred[i] == 'OBJECTIVE': objective = objective + line elif pred[i] == 'BACKGROUND': background = background + line elif pred[i] == 'METHODS': method = method + line elif pred[i] == 'RESULTS': result = result + line elif pred[i] == 'CONCLUSIONS': conclusion = conclusion + line return objective, background, method, conclusion, result def main(): st.set_page_config( page_title="SkimLit", page_icon="📄", layout="wide", initial_sidebar_state="expanded" ) st.title('SkimLit📄🔥') st.caption('An NLP model to classify abstract sentences into the role they play (e.g. objective, methods, results, etc..) to enable researchers to skim through the literature and dive deeper when necessary.') # creating model, tokenizer and labelEncoder # if PREP_MODEL: # skimlit_model, tokenizer, label_encoder = create_utils(MODEL_PATH, TOKENIZER_PATH, LABEL_ENOCDER_PATH, EMBEDDING_FILE_PATH) # PREP_MODEL = False col1, col2 = st.columns(2) with col1: st.write('#### Entre Abstract Here !!') abstract = st.text_area(label='', height=200) agree = st.checkbox('Show Example Abstract') predict = st.button('Extract !') if agree: example_input = Choose_Random_text() st.info(example_input) # make prediction button logic if predict: with col2: with st.spinner('Wait for prediction....'): skimlit_model, tokenizer, label_encoder = create_utils(MODEL_PATH, TOKENIZER_PATH, LABEL_ENOCDER_PATH, EMBEDDING_FILE_PATH) objective, background, methods, conclusion, result = model_prediction(abstract, skimlit_model, tokenizer, label_encoder) st.markdown(f'### Objective : ') st.info(objective) # st.write(f'{objective}') st.markdown(f'### Background : ') st.info(background) # st.write(f'{background}') st.markdown(f'### Methods : ') st.info(methods) # st.write(f'{methods}') st.markdown(f'### Result : ') st.info(result) # st.write(f'{result}') st.markdown(f'### Conclusion : ') st.info(conclusion) # st.write(f'{conclusion}') if __name__=='__main__': main()