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,7 +167,74 @@ def extract_text_tables_pdfplumber(pdf_file):
         else:
             print("No text extracted. The PDF might be image-based.")
             return None, None
-            
+
+# 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 = []
@@ -291,137 +348,166 @@ else:
                 extracted_data = {
                     key: (match.group(1) if match else "N/A")
                     for key, pattern in patterns.items()
-                    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.")
-    
-    # Text summarization input
-    input_text = st.text_area(
-        "Enter text to summarize",
-        height=200,
-        value=st.session_state.get("pdf_text", "") if "pdf_text" in st.session_state else ""
-    )
-    
-    if st.button("Summarize"):
-        text_to_summarize = input_text.strip()
-        if text_to_summarize:
-            try:
-                # 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 = 512
-    
-                # 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 (match := re.search(pattern, text_for_analysis, re.IGNORECASE))
+                    }
     
-                    if current_chunk:
-                        chunks.append(current_chunk.strip())
+                    # 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"]]
     
-                    return chunks
+                    # Store extracted data in a structured dictionary
+                    structured_data = {**extracted_data, "Named Entities Extracted": financial_entities}
     
-                # Function to remove duplicate sentences
-                def remove_duplicate_sentences(text):
-                    sentences = nltk.sent_tokenize(text)
-                    unique_sentences = []
-                    seen_sentences = set()
+                    # Display results
+                    st.write("Entities Found:")
+                    st.write(pd.DataFrame(financial_entities, columns=["Entity", "Label"]))
     
-                    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)
+                    st.write("Structured Data Extracted:")
+                    st.write(pd.DataFrame([structured_data]))
     
-                    return ' '.join(unique_sentences)
+            else:
+                st.error("Please provide some text for analysis.")
     
-                # 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.")
-    
-                # Summarize each 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
-    
-                if summaries:
-                    # Combine summaries
-                    combined_summary = ' '.join(summaries)
-                    # Remove duplicate sentences
-                    final_summary = remove_duplicate_sentences(combined_summary)
-                    st.write("Final Summary:")
-                    st.success(final_summary)
-                else:
-                    st.error("No summaries were generated.")
-    
-            except Exception as e:
-                st.error("An error occurred during summarization.")
-                st.text(traceback.format_exc())
-        else:
-            st.error("Please provide text to summarize.")
+        # 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",
+            height=200,
+            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:
+                try:
+                    # Display original text length
+                    input_length = len(text_to_summarize.split())
+                    st.write(f"Original text length: {input_length} words")
+                    
+                    # 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 selected chunk
+                    summaries = []
+                    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 and remove duplicates
+                        combined_summary = ' '.join(summaries)
+                        final_summary = remove_duplicate_sentences(combined_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.")
+        
+                except Exception as e:
+                    st.error("An error occurred during summarization.")
+                    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.
+            """)