Update app.py

app.py

@@ -2,16 +2,20 @@ import nltk
 nltk.data.path.append("/home/user/app/nltk_data")
 #nltk.download('stopwords')
 #nltk.download('punkt')
-import classify_abs
-import extract_abs
+from epi_pipeline import (
+    streamlit_extraction,
+    NER_Pipeline,
+    GARD_Search,
+    Classify_Pipeline
+    )
 import pandas as pd
 #pd.set_option('display.max_colwidth', None)
 import streamlit as st
 st.set_page_config(layout="wide")
-import spacy
-import tensorflow as tf
-import pickle
-import re
+#import spacy
+#import tensorflow as tf
+#import pickle
+#import re
 import plotly.graph_objects as go
 
 #### LOGO ####
@@ -58,23 +62,20 @@ filtering = st.sidebar.radio("What type of filtering would you like?",('Strict',
 extract_diseases = st.sidebar.checkbox("Extract Rare Diseases", value=False)
 
 #### MODEL LOADING ####
-
 @st.experimental_singleton(show_spinner=False)
 def load_models_experimental():
-    classify_model_vars = classify_abs.init_classify_model()
-    NER_pipeline, entity_classes = extract_abs.init_NER_pipeline()
-    GARD_dict, max_length = extract_abs.load_GARD_diseases()
-    return classify_model_vars, NER_pipeline, entity_classes, GARD_dict, max_length
+    epi_classify = Classify_Pipeline()
+    epi_extract = NER_Pipeline()
+    rd_identify = GARD_Search()
+    return epi_classify, epi_extract, rd_identify
 
 #### DOWNLOAD FUNCTION ####
-
 @st.cache
 def convert_df(df):
     # IMPORTANT: Cache the conversion to prevent computation on every rerun
     return df.to_csv().encode('utf-8')
 
 #### SANKEY FUNCTION ####
-
 #@st.cache(allow_output_mutation=True)
 @st.experimental_singleton()
 def epi_sankey(sankey_data, disease_or_gard_id):
@@ -103,9 +104,8 @@ def epi_sankey(sankey_data, disease_or_gard_id):
     return fig
 
 #### BEGIN APP ####
-
 with st.spinner('Loading Epidemiology Models and Dependencies...'):
-    classify_model_vars, NER_pipeline, entity_classes, GARD_dict, max_length = load_models_experimental()
+    epi_classify, epi_extract, rd_identify = load_models_experimental()
 loaded = st.success('All Models and Dependencies Loaded!')
 
 disease_or_gard_id = st.text_input("Input a rare disease term or NIH GARD ID.")
@@ -115,10 +115,8 @@ loaded.empty()
 st.markdown("Examples of rare diseases include [**Fellman syndrome**](https://rarediseases.info.nih.gov/diseases/1/gracile-syndrome), [**Classic Homocystinuria**](https://rarediseases.info.nih.gov/diseases/6667/classic-homocystinuria), [**7383**](https://rarediseases.info.nih.gov/diseases/7383/phenylketonuria), and [**GARD:0009941**](https://rarediseases.info.nih.gov/diseases/9941/fshmd1a). A full list of rare diseases tracked by the NIH Genetic and Rare Diseases Information Center (GARD) can be found [here](https://rarediseases.info.nih.gov/diseases/browse-by-first-letter).")
 
 if disease_or_gard_id:
-    df, sankey_data, name_gardID = extract_abs.streamlit_extraction(disease_or_gard_id, max_results, filtering,
-                                NER_pipeline, entity_classes,
-                                extract_diseases, GARD_dict, max_length,
-                                classify_model_vars)
+    df, sankey_data, name_gardID = streamlit_extraction(disease_or_gard_id, max_results, filtering,
+                                                        epi_ner, GARD_Search, extract_diseases, epi_classify)
     #IF it returns something, then continue.
     if sankey_data:
         df.replace(to_replace='None', value="None")