Update src/streamlit_app.py

src/streamlit_app.py

@@ -205,6 +205,41 @@ if 'df_ner' in st.session_state and not st.session_state.df_ner.empty:
     fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25), paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
     st.plotly_chart(fig_treemap)
 
+
+
+    dfa = pd.DataFrame(
+        data={
+            'Column Name': ['text', 'label', 'score', 'start', 'end', 'category'],
+            'Description': [
+                            'entity extracted from your text data',
+                            'label (tag) assigned to a given extracted entity',
+                            'accuracy score; how accurately a tag has been assigned to a given entity',
+                            'index of the start of the corresponding entity',
+                            'index of the end of the corresponding entity',
+                            'the broader category the entity belongs to',]}
+    )
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as myzip:
+        myzip.writestr("Summary of the results.csv", df.to_csv(index=False))
+        myzip.writestr("Glossary of tags.csv", dfa.to_csv(index=False))
+                
+    with stylable_container(
+        key="download_button",
+                    css_styles="""button { background-color: red; border: 1px solid black; padding: 5px; color: white; }""",
+    ):
+        st.download_button(
+        label="Download results and glossary (zip)",
+        data=buf.getvalue(),
+        file_name="nlpblogs_results.zip",
+        mime="application/zip",)
+                                
+    if comet_initialized:
+        experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
+        experiment.end()
+    else:
+        st.warning("No entities were found in the provided text.")
+
+
 # --- Question Answering Section ---
 @st.cache_resource
 def load_gliner_model():
@@ -280,56 +315,50 @@ if st.button("Extract Answers"):
                     st.session_state.df_qa = df_qa  # Store QA results in session state
                     st.subheader("Extracted Answers", divider="green")
                     st.dataframe(df_qa, use_container_width=True)
+                    csv_data = df_qa.to_csv(index=False).encode('utf-8')
+                    with stylable_container(
+                        key="download_button",
+                        css_styles="""button { background-color: red; border: 1px solid black; padding: 5px; color: white; }""",
+                    ):
+                        st.download_button(
+                            label="Download CSV",
+                            data=csv_data,
+                            file_name="nlpblogs_extracted_answers.csv",
+                            mime="text/csv",
+                        )
+
+                    if comet_initialized:
+                        experiment.log_metric("processing_time_seconds", elapsed_time)
+                        experiment.log_table("predicted_entities", df)
+                        experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap")
+                        
+                        experiment.end()
+                else:
+                    st.info("No answers were found in the text with the defined questions.")
+                    if comet_initialized:
+                        experiment.end()
+            except Exception as e:
+                st.error(f"An error occurred during processing: {e}")
+                st.write(f"Error details: {e}")
+                if comet_initialized:
+                    experiment.log_text(f"Error: {e}")
+                    experiment.end()
+
+
+
+
+
+
+
+
+
+
+
                 else:
                     st.warning("No answers were found for the provided questions.")
                     if 'df_qa' in st.session_state:
                         del st.session_state.df_qa
             except Exception as e:
                 st.error(f"An error occurred during answer extraction: {e}")
-                if 'df_qa' in st.session_state:
-                    del st.session_state.df_qa
-
-# --- Download Button Section ---
-def create_zip_file_and_get_bytes():
-    """Generates a zip file in memory with all available dataframes."""
-    # Define the glossary DataFrame here to ensure it's always available
-    dfa = pd.DataFrame(
-        data={
-            'Column Name': ['text', 'label', 'score', 'start', 'end', 'category'],
-            'Description': [
-                'entity extracted from your text data',
-                'label (tag) assigned to a given extracted entity',
-                'accuracy score; how accurately a tag has been assigned to a given entity',
-                'index of the start of the corresponding entity',
-                'index of the end of the corresponding entity',
-                'the broader category the entity belongs to',
-            ]
-        }
-    )
-    if 'df_ner' not in st.session_state and 'df_qa' not in st.session_state:
-        return None, None
-    buf = io.BytesIO()
-    with zipfile.ZipFile(buf, "w") as myzip:
-        if 'df_ner' in st.session_state and not st.session_state.df_ner.empty:
-            myzip.writestr("Extracted_Entities.csv", st.session_state.df_ner.to_csv(index=False))
-        if 'df_qa' in st.session_state and not st.session_state.df_qa.empty:
-            myzip.writestr("Extracted_Answers.csv", st.session_state.df_qa.to_csv(index=False))
-        myzip.writestr("Glossary_of_tags.csv", dfa.to_csv(index=False))
-    return buf.getvalue(), "nlpblogs_results.zip"
-
-st.divider()
-
-if ('df_ner' in st.session_state and not st.session_state.df_ner.empty) or \
-   ('df_qa' in st.session_state and not st.session_state.df_qa.empty):
-    zip_data, file_name = create_zip_file_and_get_bytes()
-    if zip_data:
-        with stylable_container(
-            key="download_button",
-            css_styles="""button { background-color: red; border: 1px solid black; padding: 5px; color: white; }""",
-        ):
-            st.download_button(
-                label="Download results and glossary (zip)",
-                data=zip_data,
-                file_name=file_name,
-                mime="application/zip",
-            )
+                
+