Update app.py

app.py

@@ -1,6 +1,5 @@
 import re
 import time
-import pdfplumber
 import torch
 from sentence_transformers import SentenceTransformer
 import chromadb
@@ -16,11 +15,11 @@ genai.configure(api_key="AIzaSyDXG4o4UnII5VFD1u5TaWgleG2kCfJ0Ofw")
 # Initialize Gemini model
 gemini_instance = genai.GenerativeModel('gemini-2.0-flash')
 
-# Define document paths and labels
+# Define document paths and labels (using .md files)
 doc_paths = [
-    '1002215.pdf',
-    'cancers-15-00321.pdf',
-    'ijo-57-06-1245.pdf'
+    '1002215.md',
+    'cancers-15-00321.md',
+    'ijo-57-06-1245.md'
 ]
 doc_labels = [
     "Early-stage triple negative breast cancer: the therapeutic role of immunotherapy and the prognostic value of pathological complete response",
@@ -34,27 +33,27 @@ def extract_and_chunk_docs(doc_paths, doc_labels):
     
     for doc_path, doc_label in zip(doc_paths, doc_labels):
         try:
-            # Extract text page-by-page
-            with pdfplumber.open(doc_path) as pdf:
-                full_text = []
-                for page_num, page in enumerate(pdf.pages, 1):
-                    text = page.extract_text() or ""
-                    lines = text.split('\n')
-                    for line in lines:
-                        line = line.strip()
-                        if line:
-                            full_text.append({'text': line, 'page': page_num})
+            with open(doc_path, 'r', encoding='utf-8') as md_file:
+                full_text = md_file.read()
             
-            if not full_text:
+            if not full_text.strip():
                 print(f"No content extracted from {doc_label}")
                 segmented_docs.append([])
                 doc_info.append({"label": doc_label, "gemini_structure": "No content extracted"})
                 continue
 
-            # Use Gemini to identify titles with improved prompt
-            text_for_gemini = "\n".join([entry['text'] for entry in full_text])
-            prompt = f"""You are an expert in analyzing research papers. Given the following text from a PDF, identify all section titles (e.g., Abstract, Introduction, Methods, Results, Discussion) and subsections (e.g., '2.1 Data Analysis'). Include headings that might define or characterize triple-negative breast cancer (TNBC), such as 'Definition,' 'Characteristics,' or similar, even if unconventional. Return only the titles, one per line, without explanation.
-
+            # Identify pages based on '-----'
+            pages = full_text.split('-----')
+            full_text_lines = []
+            for page_num, page_content in enumerate(pages, 1):
+                lines = page_content.split('\n')
+                for line in lines:
+                    line = line.strip()
+                    if line:
+                        full_text_lines.append({'text': line, 'page': page_num})
+            
+            text_for_gemini = "\n".join([entry['text'] for entry in full_text_lines])
+            prompt = f"""You are an expert in analyzing research papers. Given the following text from a Markdown file, identify all potential section titles (e.g., Abstract, Introduction, Methods, Results, Discussion) and subsections (e.g., '2.1 Data Analysis'). Include headings that might define or characterize triple-negative breast cancer (TNBC), such as 'Definition,' 'Characteristics,' or similar, and explicitly include 'References' as a section title if present. Only headings starting with '## **' are sections. Return only the titles (without '## **'), one per line, without explanation.
             Text:
             {text_for_gemini}
             """
@@ -81,75 +80,64 @@ def extract_and_chunk_docs(doc_paths, doc_labels):
                 doc_info.append({"label": doc_label, "gemini_structure": "Max retries exceeded"})
                 continue
             
-            # Clean and deduplicate titles
             titles = list(dict.fromkeys([title.strip() for title in titles if title.strip()]))
+            print(f"Gemini Identified Titles for {doc_label}: {titles}")
             
-            # Chunk text by titles and pages
             chunks = []
             current_chunk = ""
             current_title = "Unknown"
-            current_page = 1
             
-            for line_info in full_text:
+            for line_info in full_text_lines:
                 line = line_info['text']
-                page = line_info['page']
+                page_num = line_info['page']  # Use page number from full_text_lines
                 
                 cleaned_line = re.sub(r'(?i)copyright.*|all\s*rights\s*reserved', '', line)
                 cleaned_line = re.sub(r'\s+', ' ', cleaned_line).strip()
                 if not cleaned_line:
                     continue
                 
-                # Check if this line matches a title (case-insensitive)
-                matched_title = next((title for title in titles if cleaned_line.lower() == title.lower()), None)
+                if cleaned_line.startswith('## **'):
+                    normalized_line = re.sub(r'^##\s*\*\*|\*\*', '', cleaned_line).strip()
+                    matched_title = next((title for title in titles if normalized_line.lower() == title.lower()), None)
+                    
+                    if matched_title:
+                        # Save the previous chunk if it exists and is not "References"
+                        if current_chunk and current_title.lower() != "references":
+                            chunks.append({
+                                'text': current_chunk.strip(),
+                                'page': page_num,  # Assign the current page
+                                'section': current_title
+                            })
+                        current_title = matched_title
+                        current_chunk = ""
+                        print(f"Detected title on page {page_num}: '{current_title}'")
+                        continue
                 
-                if matched_title:
-                    # Save the previous chunk if it exists
-                    if current_chunk:
+                # Only add to chunk if current section is not "References"
+                if current_title.lower() != "references":
+                    current_chunk += " " + cleaned_line
+                    if len(current_chunk.split()) > 100:
                         chunks.append({
                             'text': current_chunk.strip(),
-                            'page': current_page,
+                            'page': page_num,  # Assign the current page
                             'section': current_title
                         })
-                    # Start a new chunk with the new title
-                    current_title = matched_title
-                    current_chunk = ""
-                    current_page = page
-                else:
-                    # Handle page breaks within the same section
-                    if page != current_page and current_chunk:
-                        chunks.append({
-                            'text': current_chunk.strip(),
-                            'page': current_page,
-                            'section': current_title
-                        })
-                        current_chunk = cleaned_line
-                        current_page = page
-                    else:
-                        # Append to current chunk, split if too long
-                        current_chunk += " " + cleaned_line
-                        if len(current_chunk.split()) > 100:
-                            chunks.append({
-                                'text': current_chunk.strip(),
-                                'page': current_page,
-                                'section': current_title
-                            })
-                            current_chunk = ""
+                        current_chunk = ""
             
-            # Save the final chunk
-            if current_chunk:
+            # Save the final chunk if it exists and is not "References"
+            if current_chunk and current_title.lower() != "references":
                 chunks.append({
                     'text': current_chunk.strip(),
-                    'page': current_page,
+                    'page': page_num,  # Assign the final page
                     'section': current_title
                 })
             
             segmented_docs.append(chunks)
-            doc_info.append({"label": doc_label, "gemini_structure": "Chunked by titles and pages"})
+            doc_info.append({"label": doc_label, "gemini_structure": "Chunked by titles starting with '## **', excluding 'References'"})
             
-            # Debugging output
             print(f"\n=== {doc_label} ===")
             print("Identified Section Titles:", titles)
-            print("\nChunks:")
+            print("\nChunks (excluding References):")
             for i, chunk in enumerate(chunks):
                 print(f"Chunk {i + 1}: [Page: {chunk['page']}, Section: '{chunk['section']}'] {chunk['text'][:100]}...")
         
@@ -159,7 +147,7 @@ def extract_and_chunk_docs(doc_paths, doc_labels):
             doc_info.append({"label": doc_label, "gemini_structure": f"Error: {str(e)}"})
     
     return segmented_docs, doc_info
-# Step 2: Generate embeddings for document segments
+
 def compute_segment_embeddings(segmented_docs):
     segment_embeddings = []
     for doc_segments in segmented_docs:
@@ -168,14 +156,13 @@ def compute_segment_embeddings(segmented_docs):
                 [seg['text'] for seg in doc_segments],
                 convert_to_tensor=False,
                 show_progress_bar=True,
-                batch_size=32  # Increased batch size for efficiency
+                batch_size=32
             )
             segment_embeddings.append(embeddings)
         else:
             segment_embeddings.append([])
     return segment_embeddings
 
-# Step 3: Store embeddings in Chroma vector store
 def save_to_vector_store(segmented_docs, segment_embeddings, doc_labels):
     db_instance = chromadb.Client()
     try:
@@ -191,20 +178,20 @@ def save_to_vector_store(segmented_docs, segment_embeddings, doc_labels):
                     documents=[segment['text']],
                     metadatas=[{
                         "label": doc_labels[i],
-                        "page": segment['page'],
+                        "page": segment['page'],  # Include page in metadata
                         "section": segment['section']
                     }],
                     ids=[f"{doc_labels[i]}seg{j}"]
                 )
     return embeddings_store
 
-# Step 4: Query the agent and generate answers
 def process_query(query, embeddings_store, doc_labels):
     query_embed = embedder.encode([query], convert_to_tensor=False)[0].tolist()
-    query_results = embeddings_store.query(query_embeddings=[query_embed], n_results=5)  # Increased to 5
+    query_results = embeddings_store.query(query_embeddings=[query_embed], n_results=3)  # Limit to top 3
     
     retrieved_contexts = []
     ref_citations = []
+    
     for doc, meta in zip(query_results["documents"][0], query_results["metadatas"][0]):
         label = meta["label"]
         page = meta["page"]
@@ -212,22 +199,13 @@ def process_query(query, embeddings_store, doc_labels):
         retrieved_contexts.append(doc)
         ref_citations.append(f"[Ref: {label}, Page: {page}, Section: '{section}']")
     
-    # Debugging: Log retrieved contexts
-    print(f"\nQuery: {query}")
-    print("Retrieved Contexts:")
-    for i, ctx in enumerate(retrieved_contexts):
-        print(f"{i + 1}: {ctx[:200]}... [Ref: {ref_citations[i]}]")
-    
     combined_context = "\n".join(retrieved_contexts) if retrieved_contexts else "No relevant context found."
     citation_str = " | ".join(ref_citations) if ref_citations else "N/A"
     
     answer_prompt = f"""You are an AI assistant for research papers. Use only the provided context to answer the query concisely (1-2 sentences max). If the context lacks a clear answer, state so briefly.
-
     Context:
     {combined_context}
-
     Query: {query}
-
     Answer:"""
     
     max_attempts = 5
@@ -247,9 +225,9 @@ def process_query(query, embeddings_store, doc_labels):
     else:
         response_text = "Error: Max retries exceeded."
     
-    return f"{response_text}\n\n*References*: {citation_str}"
+    final_response = f"{response_text}\n\n*References*: {citation_str}"
+    return final_response
 
-# Gradio chatbot function
 def chatbot_response(message, history):
     response = process_query(message, embeddings_store, doc_labels)
     return history + [{"role": "user", "content": message}, {"role": "assistant", "content": response}]
@@ -281,7 +259,7 @@ doc_segments, doc_metadata = extract_and_chunk_docs(doc_paths, doc_labels)
 segment_embeds = compute_segment_embeddings(doc_segments)
 embeddings_store = save_to_vector_store(doc_segments, segment_embeds, doc_labels)
 
-# Custom CSS (unchanged)
+# Custom CSS (simplified, no debug styling)
 css = """
 .header { text-align: center; margin-bottom: 20px; }
 .gradio-container { max-width: 900px; margin: auto; }