Update app.py

app.py

@@ -23,20 +23,43 @@ if 'summaries' not in st.session_state:
 if 'text_processor' not in st.session_state:
     st.session_state.text_processor = None
 
+def manage_resources():
+    """Clear memory and ensure resources are available"""
+    # Force garbage collection
+    gc.collect()
+    
+    # Clear CUDA cache if available
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    
+    # Set torch to use CPU
+    torch.set_num_threads(8)  # Use half of available CPU threads for each model
+
 def load_model(model_type):
-    """Load appropriate model based on type"""
+    """Load appropriate model based on type with resource management"""
+    manage_resources()
+    
     try:
+        # Set lower precision to reduce memory usage
+        torch_dtype = torch.float32
+        if torch.cuda.is_available():
+            device = "cuda"
+        else:
+            device = "cpu"
+            torch_dtype = torch.float32  # Use float32 for CPU
+        
         if model_type == "summarize":
             base_model = AutoModelForSeq2SeqLM.from_pretrained(
                 "facebook/bart-large-cnn",
-                cache_dir="./models"
+                cache_dir="./models",
+                torch_dtype=torch_dtype,
+                low_cpu_mem_usage=True
             )
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/results",
-                load_in_8bit=False,
-                device_map="auto",
-                torch_dtype=torch.float32
+                device_map=device,
+                torch_dtype=torch_dtype
             )
             tokenizer = AutoTokenizer.from_pretrained(
                 "facebook/bart-large-cnn",
@@ -45,14 +68,15 @@ def load_model(model_type):
         else:  # question_focused
             base_model = AutoModelForSeq2SeqLM.from_pretrained(
                 "GanjinZero/biobart-base",
-                cache_dir="./models"
+                cache_dir="./models",
+                torch_dtype=torch_dtype,
+                low_cpu_mem_usage=True
             )
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/biobart-finetune",
-                load_in_8bit=False,
-                device_map="auto",
-                torch_dtype=torch.float32
+                device_map=device,
+                torch_dtype=torch_dtype
             )
             tokenizer = AutoTokenizer.from_pretrained(
                 "GanjinZero/biobart-base",
@@ -148,23 +172,10 @@ def generate_focused_summary(question, abstracts, model, tokenizer):
     
     return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
+
 def main():
     st.title("🔬 Biomedical Papers Analysis")
     
-    # Sidebar
-    st.sidebar.header("About")
-    st.sidebar.info(
-        "This app analyzes biomedical research papers. Upload an Excel file "
-        "containing paper details and abstracts to:"
-        "\n- Generate individual summaries"
-        "\n- Get question-focused insights"
-    )
-    
-    # Initialize text processor if not already done
-    if st.session_state.text_processor is None:
-        with st.spinner("Loading NLP models..."):
-            st.session_state.text_processor = TextProcessor()
-    
     # File upload section
     uploaded_file = st.file_uploader(
         "Upload Excel file containing papers", 
@@ -179,74 +190,66 @@ def main():
                 df = process_excel(uploaded_file)
                 if df is not None:
                     st.session_state.processed_data = df.dropna(subset=["Abstract"])
+    
+    if st.session_state.processed_data is not None:
+        df = st.session_state.processed_data
+        st.write(f"📊 Loaded {len(df)} papers")
         
-        if st.session_state.processed_data is not None:
-            df = st.session_state.processed_data
-            st.write(f"📊 Loaded {len(df)} papers with abstracts")
-            
-            # Individual Summaries Section
-            st.header("📝 Individual Paper Summaries")
-            
-            if st.session_state.summaries is None and st.button("Generate Individual Summaries"):
-                try:
-                    with st.spinner("Generating summaries..."):
-                        # Load summarization model
+        # Individual Summaries Section
+        st.header("📝 Individual Paper Summaries")
+        
+        # Question input before the unified generate button
+        st.header("❓ Question-focused Summary (Optional)")
+        question = st.text_input("Enter your research question (optional):")
+        
+        # Unified generate button
+        if st.button("Generate Analysis"):
+            try:
+                # Step 1: Generate Individual Summaries
+                if st.session_state.summaries is None:
+                    with st.spinner("Generating individual summaries..."):
                         model, tokenizer = load_model("summarize")
                         
-                        # Process abstracts
+                        progress_text = st.empty()
                         progress_bar = st.progress(0)
-                        summaries = []
+                        summary_display = st.container()
                         
-                        for i, abstract in enumerate(df['Abstract']):
-                            summary = generate_summary(abstract, model, tokenizer)
-                            summaries.append(summary)
+                        summaries = []
+                        for i, (_, row) in enumerate(df.iterrows()):
+                            progress_text.text(f"Processing paper {i+1} of {len(df)}")
                             progress_bar.progress((i + 1) / len(df))
+                            
+                            summary = generate_summary(row['Abstract'], model, tokenizer)
+                            summaries.append(summary)
+                            
+                            with summary_display:
+                                st.write(f"**Paper {i+1}:** {row['Article Title']}")
+                                st.write(summary)
+                                st.divider()
                         
                         st.session_state.summaries = summaries
                         
-                        # Clear GPU memory
+                        # Clear memory after individual summaries
                         del model
                         del tokenizer
                         torch.cuda.empty_cache()
                         gc.collect()
                 
-                except Exception as e:
-                    st.error(f"Error generating summaries: {str(e)}")
-            
-            if st.session_state.summaries is not None:
-                # Display summaries with sorting options
-                col1, col2 = st.columns(2)
-                with col1:
-                    sort_column = st.selectbox("Sort by:", df.columns)
-                with col2:
-                    ascending = st.checkbox("Ascending order", True)
-                
-                # Create display dataframe
-                display_df = df.copy()
-                display_df['Summary'] = st.session_state.summaries
-                sorted_df = display_df.sort_values(by=sort_column, ascending=ascending)
-                
-                # Show interactive table
-                st.dataframe(sorted_df, hide_index=True)
-            
-            # Question-focused Summary Section
-            st.header("❓ Question-focused Summary")
-            question = st.text_input("Enter your research question:")
-            
-            if question and st.button("Generate Focused Summary"):
-                try:
-                    with st.spinner("Analyzing relevant papers..."):
-                        # Find relevant abstracts
+                # Step 2: Generate Question-Focused Summary (only if question is provided)
+                if question.strip():
+                    with st.spinner("Generating question-focused summary..."):
+                        # Clear memory before question processing
+                        torch.cuda.empty_cache()
+                        gc.collect()
+                        
                         results = st.session_state.text_processor.find_most_relevant_abstracts(
                             question,
                             df['Abstract'].tolist(),
                             top_k=5
                         )
                         
-                        # Load question-focused model
                         model, tokenizer = load_model("question_focused")
                         
-                        # Get relevant abstracts and generate summary
                         relevant_abstracts = df['Abstract'].iloc[results['top_indices']].tolist()
                         focused_summary = generate_focused_summary(
                             question,
@@ -255,26 +258,68 @@ def main():
                             tokenizer
                         )
                         
-                        # Display results
-                        st.subheader("Summary")
+                        st.subheader("Question-Focused Summary")
                         st.write(focused_summary)
                         
-                        # Show relevant papers
                         st.subheader("Most Relevant Papers")
                         relevant_papers = df.iloc[results['top_indices']][
                             ['Article Title', 'Authors', 'Publication Year', 'DOI']
                         ]
                         relevant_papers['Relevance Score'] = results['scores']
-                        st.dataframe(relevant_papers, hide_index=True)
+                        relevant_papers['Publication Year'] = relevant_papers['Publication Year'].astype(int)
                         
-                        # Clear GPU memory
+                        st.dataframe(
+                            relevant_papers,
+                            column_config={
+                                'Publication Year': st.column_config.NumberColumn('Year', format="%d"),
+                                'Relevance Score': st.column_config.NumberColumn('Relevance', format="%.3f")
+                            },
+                            hide_index=True
+                        )
+                        
+                        # Clear memory after question processing
                         del model
                         del tokenizer
                         torch.cuda.empty_cache()
                         gc.collect()
-                
-                except Exception as e:
-                    st.error(f"Error generating focused summary: {str(e)}")
+            
+            except Exception as e:
+                st.error(f"Error in analysis: {str(e)}")
+        
+        # Display sorted summaries if they exist
+        if st.session_state.summaries is not None:
+            st.subheader("All Paper Summaries")
+            sort_options = {
+                'Article Title': 'Article Title',
+                'Authors': 'Authors',
+                'Publication Year': 'Publication Year',
+                'Source Title': 'Source Title'
+            }
+            
+            col1, col2 = st.columns(2)
+            with col1:
+                sort_column = st.selectbox("Sort by:", list(sort_options.keys()))
+            with col2:
+                ascending = st.checkbox("Ascending order", True)
+            
+            display_df = df.copy()
+            display_df['Summary'] = st.session_state.summaries
+            display_df['Publication Year'] = display_df['Publication Year'].astype(int)
+            sorted_df = display_df.sort_values(by=sort_options[sort_column], ascending=ascending)
+            
+            st.dataframe(
+                sorted_df[['Article Title', 'Authors', 'Source Title', 
+                          'Publication Year', 'DOI', 'Summary']],
+                column_config={
+                    'Article Title': st.column_config.TextColumn('Article Title', width='medium'),
+                    'Authors': st.column_config.TextColumn('Authors', width='medium'),
+                    'Source Title': st.column_config.TextColumn('Source Title', width='medium'),
+                    'Publication Year': st.column_config.NumberColumn('Year', format="%d"),
+                    'DOI': st.column_config.TextColumn('DOI', width='small'),
+                    'Summary': st.column_config.TextColumn('Summary', width='large'),
+                },
+                hide_index=True
+            )
 
 if __name__ == "__main__":
     main()