Update app.py

app.py

@@ -9,7 +9,6 @@ import streamlit as st
 import spacy
 import tensorflow as tf
 import pickle
-import plotly.graph_objects as go
 
 ########## Title for the Web App ##########
 st.markdown('''<img src="https://huggingface.co/spaces/ncats/EpiPipeline4GARD/resolve/main/NCATS_logo.png" alt="National Center for Advancing Translational Sciences Logo" width=550>''',unsafe_allow_html=True)
@@ -70,14 +69,14 @@ def epi_sankey(sankey_data):
           pad = 15,
           thickness = 20,
           line = dict(color = "black", width = 0.5),
-          label = ["PubMed IDs Gathered", "Relevant Abstracts Gathered", "Irrelevant", "Epidemiologic", "Irrelevant"],
+          label = ["PubMed IDs Gathered", "Irrelevant Abstracts","Relevant Abstracts Gathered","Epidemiologic Abstracts","Not Epidemiologic"],
           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]
+          source = [0, 0, 2, 2],
+          target = [2, 1, 3, 4],
+          value = [relevant, gathered-relevant, epidemiologic, relevant-epidemiologic]
       ))])
     
     return fig