Update app.py

app.py

@@ -14,6 +14,7 @@ import torch
 import nltk
 from nltk.tokenize import sent_tokenize
 import traceback 
+from collections import Counter
 
 # Set Streamlit page config
 st.set_page_config(page_title="FinBrief: Financial Document Insights", layout="wide")
@@ -61,7 +62,6 @@ if summarizer is None:
 st.title("FinBrief: Financial Document Insights")
 st.write("Upload a financial document for analysis.")
 
-
 # Initialize session state
 if "nlp" not in st.session_state:
     st.session_state["nlp"] = nlp
@@ -97,16 +97,6 @@ download_nltk_punkt()
 # Debugging: Check session state initialization
 print(f"Session State - NLP: {st.session_state['nlp'] is not None}, Summarizer: {st.session_state['summarizer'] is not None}")
 
-# # Load the summarization model locally
-# try:
-#     local_model_path = "./local_models/bart-financial"
-#     summarizer = pipeline("summarization", model=local_model_path, tokenizer=local_model_path)
-#     st.write("Local summarization model loaded successfully!")
-# except Exception as e:
-#     summarizer = None  # Handle case where model is missing
-#     st.write("Failed to load local summarization model.")
-
-
 # Define regex patterns to extract structured data
 patterns = {
     "Fund Name": r"^(.*?) Fund",  # Extracts the name before "Fund"
@@ -177,8 +167,75 @@ def extract_text_tables_pdfplumber(pdf_file):
         else:
             print("No text extracted. The PDF might be image-based.")
             return None, None
-            
-def split_text_into_chunks(text, tokenizer, max_tokens=1024):
+
+# NEW: Function to evaluate chunk relevance
+def evaluate_chunk_relevance(chunk, keywords=None):
+    """
+    Evaluate the relevance of a text chunk based on various factors.
+    Returns a score representing the chunk's relevance.
+    """
+    if not keywords:
+        # Default financial keywords
+        keywords = ["fund", "portfolio", "performance", "return", "asset", "investment", 
+                    "expense", "risk", "benchmark", "allocation", "strategy", "market",
+                    "growth", "income", "dividend", "yield", "capital", "equity", "bond",
+                    "summary", "overview", "highlight", "key", "important", "significant"]
+    
+    score = 0
+    
+    # Factor 1: Length of the chunk (longer chunks often contain more information)
+    word_count = len(chunk.split())
+    score += min(word_count / 100, 5)  # Cap at 5 points
+    
+    # Factor 2: Keyword presence
+    # Count keywords in lowercase text
+    lower_chunk = chunk.lower()
+    keyword_count = sum(1 for keyword in keywords if keyword.lower() in lower_chunk)
+    keyword_density = keyword_count / max(1, word_count) * 100
+    score += min(keyword_density * 2, 10)  # Cap at 10 points
+    
+    # Factor 3: Presence of numbers (financial documents often contain important numbers)
+    number_count = len(re.findall(r'\d+\.?\d*%?', chunk))
+    score += min(number_count / 5, 5)  # Cap at 5 points
+    
+    # Factor 4: Structured information (lists, tables, etc.)
+    bullet_count = len(re.findall(r'•|\*|-|–|[0-9]+\.', chunk))
+    score += min(bullet_count, 5)  # Cap at 5 points
+    
+    # Factor 5: Presence of section headers
+    header_patterns = [
+        r'^[A-Z][A-Za-z\s]+:',  # Title followed by colon
+        r'^[A-Z][A-Z\s]+',      # ALL CAPS text
+        r'^\d+\.\s+[A-Z]'       # Numbered section
+    ]
+    header_count = sum(1 for pattern in header_patterns if re.search(pattern, chunk, re.MULTILINE))
+    score += min(header_count * 2, 5)  # Cap at 5 points
+    
+    return score
+
+# NEW: Function to rank and select the best chunks
+def rank_and_select_chunks(chunks, max_chunks=5, keywords=None):
+    """
+    Rank chunks by relevance and return the top chunks.
+    """
+    # Evaluate each chunk
+    chunk_scores = [(chunk, evaluate_chunk_relevance(chunk, keywords)) for chunk in chunks]
+    
+    # Sort chunks by score (highest first)
+    sorted_chunks = sorted(chunk_scores, key=lambda x: x[1], reverse=True)
+    
+    # Select the top N chunks
+    top_chunks = [chunk for chunk, score in sorted_chunks[:max_chunks]]
+    
+    # Print scores for debugging
+    print("Chunk scores:")
+    for i, (chunk, score) in enumerate(sorted_chunks):
+        print(f"Chunk {i+1}: Score {score:.2f}, Length {len(chunk.split())} words")
+        print(f"First 100 chars: {chunk[:100]}...")
+    
+    return top_chunks
+
+def split_text_into_chunks(text, tokenizer, max_tokens=512):
     sentences = nltk.sent_tokenize(text)
     chunks = []
     current_chunk = ''
@@ -269,7 +326,7 @@ pdf_tables = st.session_state["pdf_tables"]
 
 # Ensure that the models are loaded
 if nlp is None or summarizer is None:
-    st.error("Models are not properly loaded. Please check your model paths and installation.")
+    st.error("Models are not properly loaded. Please check model paths and installation.")
 else:
     # Step 3: Named Entity Recognition (NER)
     st.subheader("NER Analysis")
@@ -284,7 +341,7 @@ else:
     if st.button("Analyze"):
         # Ensure full extracted text is used for analysis
         text_for_analysis = st.session_state["pdf_text"].strip() if st.session_state["pdf_text"] else example_text.strip()
-
+    
         if text_for_analysis:
             with st.spinner("Analyzing text..."):
                 # Extract structured financial data using regex (Now using full text)
@@ -294,27 +351,43 @@ else:
                     if (match := re.search(pattern, text_for_analysis, re.IGNORECASE))
                 }
 
-                # Use spaCy to extract additional financial terms (Now using full text)
                 doc = nlp(text_for_analysis)
                 financial_entities = [(ent.text, ent.label_) for ent in doc.ents if ent.label_ in ["MONEY", "PERCENT", "ORG", "DATE"]]
-
+    
                 # Store extracted data in a structured dictionary
                 structured_data = {**extracted_data, "Named Entities Extracted": financial_entities}
-
+    
                 # Display results
                 st.write("Entities Found:")
                 st.write(pd.DataFrame(financial_entities, columns=["Entity", "Label"]))
-
+    
                 st.write("Structured Data Extracted:")
                 st.write(pd.DataFrame([structured_data]))
-
+    
         else:
             st.error("Please provide some text for analysis.")
-
+    
     # Step 4: Summarization
     st.subheader("Summarization")
     st.write("Generate concise summaries of financial documents.")
     
+    # Add customization options for summarization with chunk selection
+    st.sidebar.header("Summarization Settings")
+    max_chunks_to_process = st.sidebar.slider(
+        "Max chunks to summarize", 
+        min_value=1, 
+        max_value=10, 
+        value=3,
+        help="Select fewer chunks for faster processing but less comprehensive summaries"
+    )
+    
+    # Allow users to add custom keywords
+    custom_keywords = st.sidebar.text_input(
+        "Add custom keywords (comma separated)",
+        value="",
+        help="Add domain-specific keywords to improve chunk selection"
+    )
+    
     # Text summarization input
     input_text = st.text_area(
         "Enter text to summarize",
@@ -322,6 +395,9 @@ else:
         value=st.session_state.get("pdf_text", "") if "pdf_text" in st.session_state else ""
     )
     
+    # Add option to see chunk selection details
+    show_chunk_details = st.sidebar.checkbox("Show chunk selection details", value=False)
+    
     if st.button("Summarize"):
         text_to_summarize = input_text.strip()
         if text_to_summarize:
@@ -329,94 +405,88 @@ else:
                 # Display original text length
                 input_length = len(text_to_summarize.split())
                 st.write(f"Original text length: {input_length} words")
-    
-                # Define the maximum number of tokens the model can handle
-                max_input_tokens = 1024  # BART's maximum input length
-    
-                # Function to split text into chunks based on tokens (modified to avoid overlaps)
-                def split_text_into_chunks(text, tokenizer, max_tokens=max_input_tokens):
-                    sentences = nltk.sent_tokenize(text)
-                    chunks = []
-                    current_chunk = ''
-                    current_length = 0
-    
-                    for sentence in sentences:
-                        sentence_tokens = tokenizer.encode(sentence, add_special_tokens=False)
-                        sentence_length = len(sentence_tokens)
-    
-                        # If adding the sentence exceeds max_tokens, start a new chunk
-                        if current_length + sentence_length > max_tokens:
-                            if current_chunk:
-                                chunks.append(current_chunk.strip())
-                            current_chunk = sentence
-                            current_length = sentence_length
-                        else:
-                            current_chunk += ' ' + sentence
-                            current_length += sentence_length
-    
-                    if current_chunk:
-                        chunks.append(current_chunk.strip())
-    
-                    return chunks
-    
-                # Function to remove duplicate sentences
-                def remove_duplicate_sentences(text):
-                    sentences = nltk.sent_tokenize(text)
-                    unique_sentences = []
-                    seen_sentences = set()
-    
-                    for sentence in sentences:
-                        normalized_sentence = sentence.strip().lower()
-                        if normalized_sentence not in seen_sentences:
-                            seen_sentences.add(normalized_sentence)
-                            unique_sentences.append(sentence)
-    
-                    return ' '.join(unique_sentences)
+                
+                # Process custom keywords if provided
+                keywords = None
+                if custom_keywords:
+                    keywords = [kw.strip() for kw in custom_keywords.split(",") if kw.strip()]
+                    st.write(f"Using custom keywords: {', '.join(keywords)}")
     
                 # Split the text into manageable chunks
                 chunks = split_text_into_chunks(text_to_summarize, tokenizer)
                 st.write(f"Text has been split into {len(chunks)} chunks.")
+                
+                # NEW: Rank and select the best chunks instead of processing all of them
+                selected_chunks = rank_and_select_chunks(
+                    chunks, 
+                    max_chunks=max_chunks_to_process,
+                    keywords=keywords
+                )
+                
+                st.write(f"Selected {len(selected_chunks)} highest-ranked chunks for summarization.")
+                
+                # Show chunk selection details if requested
+                if show_chunk_details:
+                    with st.expander("Chunk Selection Details"):
+                        for i, chunk in enumerate(selected_chunks):
+                            st.markdown(f"**Chunk {i+1}**")
+                            st.write(f"Length: {len(chunk.split())} words")
+                            st.text(chunk[:300] + "..." if len(chunk) > 300 else chunk)
+                            st.write("---")
     
-                # Summarize each chunk
+                # Summarize each selected chunk
                 summaries = []
-                for i, chunk in enumerate(chunks):
-                    st.write(f"Summarizing chunk {i+1}/{len(chunks)}...")
-                    # Adjust summary length parameters as needed
-                    chunk_length = len(chunk.split())
-                    max_summary_length = min(150, chunk_length // 2)
-                    min_summary_length = max(50, max_summary_length // 2)
-    
-                    try:
-                        summary_output = summarizer(
-                            chunk,
-                            max_length=max_summary_length,
-                            min_length=min_summary_length,
-                            do_sample=False,
-                            truncation=True
-                        )
-                        chunk_summary = summary_output[0]['summary_text'].strip()
-    
-                        if not chunk_summary:
-                            st.warning(f"The summary for chunk {i+1} is empty.")
-                        else:
-                            summaries.append(chunk_summary)
-                            # Optionally display the summary of the current chunk
-                            # st.write(f"Summary of chunk {i+1}:")
-                            # st.write(chunk_summary)
-                            # st.write("---")
-    
-                    except Exception as e:
-                        st.error(f"Summarization failed for chunk {i+1}: {e}")
-                        st.text(traceback.format_exc())
-                        continue
+                with st.spinner(f"Summarizing {len(selected_chunks)} chunks..."):
+                    for i, chunk in enumerate(selected_chunks):
+                        st.write(f"Summarizing chunk {i+1}/{len(selected_chunks)}...")
+                        # Adjust summary length parameters as needed
+                        chunk_length = len(chunk.split())
+                        max_summary_length = min(150, chunk_length // 2)
+                        min_summary_length = max(50, max_summary_length // 2)
+        
+                        try:
+                            summary_output = summarizer(
+                                chunk,
+                                max_length=max_summary_length,
+                                min_length=min_summary_length,
+                                do_sample=False,
+                                truncation=True
+                            )
+                            chunk_summary = summary_output[0]['summary_text'].strip()
+        
+                            if not chunk_summary:
+                                st.warning(f"The summary for chunk {i+1} is empty.")
+                            else:
+                                summaries.append(chunk_summary)
+        
+                        except Exception as e:
+                            st.error(f"Summarization failed for chunk {i+1}: {e}")
+                            st.text(traceback.format_exc())
+                            continue
     
                 if summaries:
-                    # Combine summaries
+                    # Combine summaries and remove duplicates
                     combined_summary = ' '.join(summaries)
-                    # Remove duplicate sentences
                     final_summary = remove_duplicate_sentences(combined_summary)
-                    st.write("Final Summary:")
+                    
+                    # Calculate compression ratio
+                    summary_length = len(final_summary.split())
+                    compression_ratio = (1 - summary_length / input_length) * 100
+                    
+                    st.subheader("Final Summary")
                     st.success(final_summary)
+                    st.write(f"Summary length: {summary_length} words ({compression_ratio:.1f}% compression)")
+                    
+                    # Display summary statistics
+                    st.subheader("Summary Statistics")
+                    stats_col1, stats_col2 = st.columns(2)
+                    with stats_col1:
+                        st.metric("Original Length", f"{input_length} words")
+                        st.metric("Total Chunks", str(len(chunks)))
+                    with stats_col2:
+                        st.metric("Summary Length", f"{summary_length} words")
+                        st.metric("Chunks Processed", str(len(selected_chunks)))
+                    
                 else:
                     st.error("No summaries were generated.")
     
@@ -425,3 +495,18 @@ else:
                 st.text(traceback.format_exc())
         else:
             st.error("Please provide text to summarize.")
+            
+    # Add help information
+    st.sidebar.markdown("---")
+    with st.sidebar.expander("How Chunk Selection Works"):
+        st.markdown("""
+        The chunk selection algorithm ranks text chunks based on:
+        
+        1. **Keyword density** - Presence of financial terms
+        2. **Length** - Longer chunks often contain more information
+        3. **Numbers** - Financial documents with numbers are often important
+        4. **Structure** - Lists and bullet points signal key information
+        5. **Headers** - Section headers often introduce important content
+        
+        Adjust the settings above to customize the selection process.
+        """)