Update app.py

app.py

@@ -62,6 +62,24 @@ def load_models():
     GARD_dict, max_length = extract_abs.load_GARD_diseases()
     return classify_tokenizer, classify_model, NER_pipeline, entity_classes, GARD_dict, max_length
 
+def epi_sankey(sankey_data):
+    gathered, relevant, epidemiologic = sankey_data
+    fig = go.Figure(data=[go.Sankey(
+    node = dict(
+      pad = 15,
+      thickness = 20,
+      line = dict(color = "black", width = 0.5),
+      label = ["PubMed IDs Gathered", "Relevant Abstracts Gathered", "Irrelevant", "Epidemiologic", "Irrelevant"],
+      color = "blue"
+    ),
+    #label = ["A1", "A2", "B1", "B2", "C1", "C2"]
+    link = dict(
+      source = [0, 0, 0, 1, 3]
+      target = [0, 1, 2, 3, 4],
+      value = [gathered, gathered-relevant, relevant, epidemiologic, gathered-epidemiologic]
+  ))])
+  return fig
+
 with st.spinner('Loading Epidemiology Models and Dependencies...'):
     classify_model_vars, NER_pipeline, entity_classes, GARD_dict, max_length = load_models_experimental()
     #classify_tokenizer, classify_model, NER_pipeline, entity_classes, GARD_dict, max_length = load_models()
@@ -84,7 +102,6 @@ if disease_or_gard_id:
                                 extract_diseases,GARD_dict, max_length, 
                                 classify_model_vars)
     st.dataframe(df, height=100)
-    csv = convert_df(df)
     st.download_button(
         label="Download epidemiology results for "+disease_or_gard_id+" as CSV",
         data=df.to_csv().encode('utf-8'),
@@ -95,5 +112,7 @@ if disease_or_gard_id:
     gathered, relevant, epidemiologic = sankey_data
     if st.button('Display Sankey Diagram'):
         st.write('Sankey Diagram here')
+        #fig = epi_sankey(sankey_data)
+        #st.plotly_chart(fig, use_container_width=True)
         pass
 # st.code(body, language="python")