Update app.py

app.py

@@ -1,11 +1,10 @@
 import streamlit as st
 import pandas as pd
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLm
 from peft import PeftModel
 from text_processing import TextProcessor
 import gc
-import time
 from pathlib import Path
 
 # Configure page
@@ -22,40 +21,31 @@ if 'summaries' not in st.session_state:
     st.session_state.summaries = None
 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
+if 'processing_started' not in st.session_state:
+    st.session_state.processing_started = False
+if 'focused_summary_generated' not in st.session_state:
+    st.session_state.focused_summary_generated = False
 
 def load_model(model_type):
-    """Load appropriate model based on type with resource management"""
-    manage_resources()
-    
+    """Load appropriate model based on type with proper memory management"""
     try:
+        # Clear any existing cached data
+        torch.cuda.empty_cache()
+        gc.collect()
+        
         if model_type == "summarize":
             base_model = AutoModelForSeq2SeqLM.from_pretrained(
                 "facebook/bart-large-cnn",
                 cache_dir="./models",
-                device_map=None,  # Explicitly set to None for CPU
+                low_cpu_mem_usage=True,
                 torch_dtype=torch.float32
-            ).to("cpu")  # Force CPU
-            
+            )
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/results",
-                device_map=None,  # Explicitly set to None for CPU
-                torch_dtype=torch.float32,
-                is_trainable=False  # Set to inference mode
-            ).to("cpu")  # Force CPU
-            
+                device_map="auto",
+                torch_dtype=torch.float32
+            )
             tokenizer = AutoTokenizer.from_pretrained(
                 "facebook/bart-large-cnn",
                 cache_dir="./models"
@@ -64,36 +54,43 @@ def load_model(model_type):
             base_model = AutoModelForSeq2SeqLM.from_pretrained(
                 "GanjinZero/biobart-base",
                 cache_dir="./models",
-                device_map=None,  # Explicitly set to None for CPU
+                low_cpu_mem_usage=True,
                 torch_dtype=torch.float32
-            ).to("cpu")  # Force CPU
-            
+            )
             model = PeftModel.from_pretrained(
                 base_model, 
                 "pendar02/biobart-finetune",
-                device_map=None,  # Explicitly set to None for CPU
-                torch_dtype=torch.float32,
-                is_trainable=False  # Set to inference mode
-            ).to("cpu")  # Force CPU
-            
+                device_map="auto",
+                torch_dtype=torch.float32
+            )
             tokenizer = AutoTokenizer.from_pretrained(
                 "GanjinZero/biobart-base",
                 cache_dir="./models"
             )
         
-        model.eval()  # Set to evaluation mode
+        model.eval()
         return model, tokenizer
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")
         raise
 
+def cleanup_model(model, tokenizer):
+    """Properly cleanup model resources"""
+    try:
+        del model
+        del tokenizer
+        torch.cuda.empty_cache()
+        gc.collect()
+    except Exception:
+        pass
+
 @st.cache_data
 def process_excel(uploaded_file):
     """Process uploaded Excel file"""
     try:
         df = pd.read_excel(uploaded_file)
         required_columns = ['Abstract', 'Article Title', 'Authors', 
-                            'Source Title', 'Publication Year', 'DOI']
+                          'Source Title', 'Publication Year', 'DOI']
         
         # Check required columns
         missing_columns = [col for col in required_columns if col not in df.columns]
@@ -127,9 +124,18 @@ def generate_summary(text, model, tokenizer):
     if not isinstance(text, str) or not text.strip():
         return "No abstract available to summarize."
     
+    # Check if abstract is too short
+    word_count = len(text.split())
+    if word_count < 50:  # Threshold for "short" abstracts
+        return text  # Return original text for very short abstracts
+    
     # Preprocess the text first
     formatted_text = preprocess_text(text)
     
+    # Adjust generation parameters based on input length
+    max_length = min(150, word_count + 50)  # Dynamic max length
+    min_length = min(50, word_count)  # Dynamic min length
+    
     inputs = tokenizer(formatted_text, return_tensors="pt", max_length=1024, truncation=True)
     
     with torch.no_grad():
@@ -137,15 +143,22 @@ def generate_summary(text, model, tokenizer):
             **{
                 "input_ids": inputs["input_ids"],
                 "attention_mask": inputs["attention_mask"],
-                "max_length": 150,
-                "min_length": 50,
+                "max_length": max_length,
+                "min_length": min_length,
                 "num_beams": 4,
                 "length_penalty": 2.0,
-                "early_stopping": True
+                "early_stopping": True,
+                "no_repeat_ngram_size": 3  # Prevent repetition of phrases
             }
         )
     
-    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    
+    # Post-process summary
+    if summary.lower() == text.lower() or len(summary.split()) / word_count > 0.9:
+        return text  # Return original if summary is too similar
+        
+    return summary
 
 def generate_focused_summary(question, abstracts, model, tokenizer):
     """Generate focused summary based on question"""
@@ -173,11 +186,6 @@ def generate_focused_summary(question, abstracts, model, tokenizer):
 def main():
     st.title("🔬 Biomedical Papers Analysis")
     
-    # 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", 
@@ -185,6 +193,10 @@ def main():
         help="File must contain: Abstract, Article Title, Authors, Source Title, Publication Year, DOI"
     )
     
+    # Question input - moved up but hidden initially
+    question_container = st.empty()
+    question = ""
+    
     if uploaded_file is not None:
         # Process Excel file
         if st.session_state.processed_data is None:
@@ -192,146 +204,144 @@ 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")
         
-        # 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.processed_data is not None:
+            df = st.session_state.processed_data
+            st.write(f"📊 Loaded {len(df)} papers with abstracts")
+            
+            # Get question before processing
+            with question_container:
+                question = st.text_input(
+                    "Enter your research question (optional):",
+                    help="If provided, a question-focused summary will be generated after individual summaries"
+                )
+            
+            # Single button for both processes
+            if not st.session_state.get('processing_started', False):
+                if st.button("Start Analysis"):
+                    st.session_state.processing_started = True
+            
+            # Show processing status and results
+            if st.session_state.get('processing_started', False):
+                # Individual Summaries Section
+                st.header("📝 Individual Paper Summaries")
+                
                 if st.session_state.summaries is None:
-                    with st.spinner("Generating individual summaries..."):
-                        model, tokenizer = load_model("summarize")
-                        
-                        progress_text = st.empty()
-                        progress_bar = st.progress(0)
-                        
-                        # Create a table for live updates
-                        summary_table = st.empty()
-                        summaries = []
-                        table_data = []
-                        
-                        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))
+                    try:
+                        with st.spinner("Generating summaries..."):
+                            # Load summarization model
+                            model, tokenizer = load_model("summarize")
                             
-                            summary = generate_summary(row['Abstract'], model, tokenizer)
-                            summaries.append(summary)
+                            # Process abstracts with real-time updates
+                            summaries = []
+                            progress_bar = st.progress(0)
+                            summary_display = st.empty()
                             
-                            # Update table data
-                            table_data.append({
-                                "PAPER": f"{row['Article Title']}\n{row['Authors']}\nDOI: {row['DOI']}",
-                                "SUMMARY": summary
-                            })
-                            summary_table.dataframe(
-                                pd.DataFrame(table_data),
-                                column_config={
-                                    "PAPER": st.column_config.TextColumn("PAPER", width=300),
-                                    "SUMMARY": st.column_config.TextColumn("SUMMARY", width="medium")
-                                },
-                                hide_index=True
-                            )
-                        
-                        st.session_state.summaries = summaries
-                        
-                        # Clear memory after individual summaries
-                        del model
-                        del tokenizer
-                        torch.cuda.empty_cache()
-                        gc.collect()
+                            for i, (_, row) in enumerate(df.iterrows()):
+                                summary = generate_summary(row['Abstract'], model, tokenizer)
+                                summaries.append(summary)
+                                
+                                # Update progress and show current summary
+                                progress = (i + 1) / len(df)
+                                progress_bar.progress(progress)
+                                summary_display.write(f"Processing paper {i+1}/{len(df)}:\n{row['Article Title']}")
+                                
+                            st.session_state.summaries = summaries
+                            
+                            # Cleanup first model
+                            cleanup_model(model, tokenizer)
+                    
+                    except Exception as e:
+                        st.error(f"Error generating summaries: {str(e)}")
                 
-                # Step 2: Generate Question-Focused Summary (only if question is provided)
+                # Display summaries with improved sorting
+                if st.session_state.summaries is not None:
+                    col1, col2 = st.columns(2)
+                    with col1:
+                        sort_options = ['Article Title', 'Authors', 'Publication Year', 'Source Title']
+                        sort_column = st.selectbox("Sort by:", sort_options)
+                    with col2:
+                        ascending = st.checkbox("Ascending order", True)
+                    
+                    # Create display dataframe with formatted year
+                    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_column, ascending=ascending)
+                    
+                    # Apply custom formatting
+                    st.markdown("""
+                    <style>
+                    .stDataFrame {
+                        font-size: 16px;
+                    }
+                    .stDataFrame td {
+                        white-space: normal !important;
+                        padding: 8px !important;
+                    }
+                    </style>
+                    """, unsafe_allow_html=True)
+                    
+                    st.dataframe(
+                        sorted_df[['Article Title', 'Authors', 'Source Title', 
+                                 'Publication Year', 'DOI', 'Summary']],
+                        hide_index=True
+                    )
+                
+                # Question-focused Summary Section (only if question provided)
                 if question.strip():
-                    with st.spinner("Generating question-focused summary..."):
-                        # Clear memory before question processing
-                        torch.cuda.empty_cache()
-                        gc.collect()
-                        
-                        # Find relevant abstracts
-                        results = st.session_state.text_processor.find_most_relevant_abstracts(
-                            question,
-                            df['Abstract'].tolist(),
-                            top_k=5
-                        )
-                        
-                        # Load question model
-                        model, tokenizer = load_model("question_focused")
-                        
-                        relevant_abstracts = df['Abstract'].iloc[results['top_indices']].tolist()
-                        focused_summary = generate_focused_summary(
-                            question,
-                            relevant_abstracts,
-                            model,
-                            tokenizer
-                        )
+                    st.header("❓ Question-focused Summary")
+                    
+                    if not st.session_state.get('focused_summary_generated', False):
+                        try:
+                            with st.spinner("Analyzing relevant papers..."):
+                                # Initialize text processor if needed
+                                if st.session_state.text_processor is None:
+                                    st.session_state.text_processor = TextProcessor()
+                                
+                                # Find relevant abstracts
+                                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")
+                                
+                                # Generate focused summary
+                                relevant_abstracts = df['Abstract'].iloc[results['top_indices']].tolist()
+                                focused_summary = generate_focused_summary(
+                                    question,
+                                    relevant_abstracts,
+                                    model,
+                                    tokenizer
+                                )
+                                
+                                # Store results
+                                st.session_state.focused_summary = focused_summary
+                                st.session_state.relevant_papers = df.iloc[results['top_indices']]
+                                st.session_state.relevance_scores = results['scores']
+                                st.session_state.focused_summary_generated = True
+                                
+                                # Cleanup second model
+                                cleanup_model(model, tokenizer)
                         
-                        st.subheader("Question-Focused Summary")
-                        st.write(focused_summary)
+                        except Exception as e:
+                            st.error(f"Error generating focused summary: {str(e)}")
+                    
+                    # Display focused summary results
+                    if st.session_state.get('focused_summary_generated', False):
+                        st.subheader("Summary")
+                        st.write(st.session_state.focused_summary)
                         
                         st.subheader("Most Relevant Papers")
-                        relevant_papers = df.iloc[results['top_indices']][
+                        relevant_papers = st.session_state.relevant_papers[
                             ['Article Title', 'Authors', 'Publication Year', 'DOI']
-                        ]
-                        relevant_papers['Relevance Score'] = results['scores']
+                        ].copy()
+                        relevant_papers['Relevance Score'] = st.session_state.relevance_scores
                         relevant_papers['Publication Year'] = relevant_papers['Publication Year'].astype(int)
-                        
-                        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 in analysis: {str(e)}")
-        
-        # Display sorted summaries if they exist
-        if st.session_state.summaries is not None:
-            st.header("📝 Individual Paper Summaries")
-            col1, col2 = st.columns([2, 1])
-            with col1:
-                sort_by = st.selectbox(
-                    "Sort By",
-                    ["Article Title", "Publication Year"],
-                    key="sort_summaries"
-                )
-            with col2:
-                ascending = st.checkbox("Ascending order", True, key="sort_order")
-            
-            # Create display dataframe
-            display_df = df.copy()
-            display_df['PAPER'] = display_df.apply(
-                lambda x: f"{x['Article Title']}\n{x['Authors']}\nDOI: {x['DOI']}", 
-                axis=1
-            )
-            display_df['SUMMARY'] = st.session_state.summaries
-            
-            # Sort the dataframe
-            sorted_df = display_df.sort_values(by=sort_by, ascending=ascending)
-            
-            # Display the table
-            st.dataframe(
-                sorted_df[['PAPER', 'SUMMARY']],
-                column_config={
-                    "PAPER": st.column_config.TextColumn("PAPER", width=300),
-                    "SUMMARY": st.column_config.TextColumn("SUMMARY", width="medium")
-                },
-                hide_index=True
-            )
+                        st.dataframe(relevant_papers, hide_index=True)
 
 if __name__ == "__main__":
     main()