Update app.py

app.py

@@ -1,3 +1,309 @@
+# # ---------------------------------------------------------------------------------------
+# # Imports and Options
+# # ---------------------------------------------------------------------------------------
+# import streamlit as st
+# import pandas as pd
+# import requests
+# import re
+# import fitz  # PyMuPDF
+# import io
+# import matplotlib.pyplot as plt
+# from PIL import Image
+# from mlx_vlm import load, generate
+# from mlx_vlm.prompt_utils import apply_chat_template
+# from mlx_vlm.utils import load_config, stream_generate
+# from docling_core.types.doc.document import DocTagsDocument, DoclingDocument
+
+# # Set Streamlit to wide mode
+# # st.set_page_config(layout="wide")
+
+# # ---------------------------------------------------------------------------------------
+# # API Configuration
+# # ---------------------------------------------------------------------------------------
+# API_URL = "https://api.stack-ai.com/inference/v0/run/2df89a6c-a4af-4576-880e-27058e498f02/67acad8b0603ba4631db38e7"
+# headers = {
+#     'Authorization': 'Bearer a9e4979e-cdbe-49ea-a193-53562a784805',
+#     'Content-Type': 'application/json'
+# }
+
+# # ---------------------------------------------------------------------------------------
+# # Survey Analysis Class
+# # ---------------------------------------------------------------------------------------
+# class SurveyAnalysis:
+#     def __init__(self, api_key=None):
+#         self.api_key = api_key
+
+#     def prepare_llm_input(self, survey_response, topics):
+#         # Create topic description string from user input
+#         topic_descriptions = "\n".join([f"- **{topic}**: {description}" for topic, description in topics.items()])
+        
+#         llm_input = f"""
+# Your task is to review PDF docling and extract information related to the provided topics. Here are the topic descriptions:
+
+# {topic_descriptions}
+
+# **Instructions:**
+# - Extract and summarize the PDF focusing only on the provided topics.
+# - If a topic is not mentioned in the notes, it should not be included in the Topic_Summary.
+# - Use **exact quotes** from the original text for each point in your Topic_Summary.
+# - Exclude erroneous content.
+# - Do not add additional explanations or instructions.
+
+# **Format your response as follows:**
+# [Topic]
+# - "Exact quote"
+# - "Exact quote"
+# - "Exact quote"
+
+# **Meeting Notes:**
+# {survey_response}
+# """
+#         return llm_input
+
+#     def query_api(self, payload):
+#         response = requests.post(API_URL, headers=headers, json=payload)
+#         return response.json()
+
+#     def extract_meeting_notes(self, response):
+#         output = response.get('outputs', {}).get('out-0', '')
+#         return output
+
+#     def process_dataframe(self, df, topics):
+#         results = []
+#         for _, row in df.iterrows():
+#             llm_input = self.prepare_llm_input(row['Document_Text'], topics)
+#             payload = {
+#                 "user_id": "<USER or Conversation ID>",  
+#                 "in-0": llm_input
+#             }
+#             response = self.query_api(payload)
+#             meeting_notes = self.extract_meeting_notes(response)
+#             results.append({
+#                 'Document_Text': row['Document_Text'],
+#                 'Topic_Summary': meeting_notes
+#             })
+        
+#         result_df = pd.DataFrame(results)
+#         df = df.reset_index(drop=True)
+#         return pd.concat([df, result_df[['Topic_Summary']]], axis=1)
+
+# # ---------------------------------------------------------------------------------------
+# # Function to Extract Excerpts
+# # ---------------------------------------------------------------------------------------
+# def extract_excerpts(processed_df):
+#     new_rows = []
+    
+#     for _, row in processed_df.iterrows():
+#         Topic_Summary = row['Topic_Summary']
+        
+#         # Split the Topic_Summary by topic
+#         sections = re.split(r'\n(?=\[)', Topic_Summary)
+        
+#         for section in sections:
+#             # Extract the topic
+#             topic_match = re.match(r'\[([^\]]+)\]', section)
+#             if topic_match:
+#                 topic = topic_match.group(1)
+                
+#                 # Extract all excerpts within the section
+#                 excerpts = re.findall(r'- "([^"]+)"', section)
+                
+#                 for excerpt in excerpts:
+#                     new_rows.append({
+#                         'Document_Text': row['Document_Text'],
+#                         'Topic_Summary': row['Topic_Summary'],
+#                         'Excerpt': excerpt,
+#                         'Topic': topic
+#                     })
+    
+#     return pd.DataFrame(new_rows)
+
+# #------------------------------------------------------------------------
+# # Streamlit Configuration
+# #------------------------------------------------------------------------
+
+# # Set page configuration
+# st.set_page_config(
+#     page_title="Choose Your Own Adventure (Topic Extraction) PDF Analysis App", 
+#     page_icon=":bar_chart:", 
+#     layout="centered", 
+#     initial_sidebar_state="auto",
+#     menu_items={
+#         'Get Help': 'mailto:support@mtss.ai',
+#         'About': "This app is built to support PDF analysis"
+#     }
+# )
+
+# #------------------------------------------------------------------------
+# # Sidebar
+# #------------------------------------------------------------------------
+
+# # Sidebar with image
+# with st.sidebar:
+#     # Set the desired width in pixels
+#     image_width = 300  
+#     # Define the path to the image
+#     # image_path = "steelcase_small.png"
+#     image_path = "mtss.ai_small.png"
+#     # Display the image
+#     st.image(image_path, width=image_width)
+
+#     # Additional sidebar content
+
+#     with st.expander("**MTSS.ai**", expanded=True):
+#         st.write("""
+#         - **Support**: Cheyne LeVesseur PhD  
+#         - **Email**: support@mtss.ai
+#         """)
+#     st.divider()
+#     st.subheader('Instructions')
+
+#     Instructions = """
+#     - **Step 1**: Upload your PDF file.
+#     - **Step 2**: Review the processed text.
+#     - **Step 3**: Add your topics and descriptions of interest.
+#     - **Step 4**: Review the extracted excerpts and classifications, and topic distribution and frequency.
+#     - **Step 5**: Review bar charts of topics.
+#     - **Step 6**: Download the processed data as a CSV file.
+#     """
+#     st.markdown(Instructions)
+
+# # Load SmolDocling model ()
+# @st.cache_resource
+# def load_smol_docling():
+#     model_path = "ds4sd/SmolDocling-256M-preview"
+#     model, processor = load(model_path)
+#     config = load_config(model_path)
+#     return model, processor, config
+
+# model, processor, config = load_smol_docling()
+
+# # Convert PDF to images
+# def convert_pdf_to_images(pdf_file):
+#     images = []
+#     doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
+#     for page_number in range(len(doc)):
+#         page = doc.load_page(page_number)
+#         pix = page.get_pixmap(dpi=300)  # Higher DPI for clarity
+#         img_data = pix.tobytes("png")
+#         image = Image.open(io.BytesIO(img_data))
+#         images.append(image)
+#     return images
+
+# # Extract structured markdown text using SmolDocling (mlx_vlm)
+# def extract_markdown_from_image(image):
+#     prompt = "Convert this page to docling."
+#     formatted_prompt = apply_chat_template(processor, config, prompt, num_images=1)
+#     output = ""
+
+#     for token in stream_generate(
+#         model, processor, formatted_prompt, [image], max_tokens=4096, verbose=False):
+#         output += token.text
+#         if "</doctag>" in token.text:
+#             break
+
+#     # Convert DocTags to Markdown
+#     doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([output], [image])
+#     doc = DoclingDocument(name="ExtractedDocument")
+#     doc.load_from_doctags(doctags_doc)
+#     markdown_text = doc.export_to_markdown()
+#     return markdown_text
+
+# # Streamlit UI
+# st.title("Choose Your Own Adventure (Topic Extraction) PDF Analysis App")
+
+# uploaded_file = st.file_uploader("Upload PDF file", type=["pdf"])
+
+# if uploaded_file:
+#     with st.spinner("Processing PDF..."):
+#         images = convert_pdf_to_images(uploaded_file)
+
+#         markdown_texts = []
+#         for idx, image in enumerate(images):
+#             markdown_text = extract_markdown_from_image(image)
+#             markdown_texts.append(markdown_text)
+
+#         df = pd.DataFrame({'Document_Text': markdown_texts})
+
+#     st.success("PDF processed successfully!")
+
+#     # Check if extraction was successful
+#     if df.empty or df['Document_Text'].isnull().all():
+#         st.error("No meaningful text extracted from the PDF.")
+#         st.stop()
+
+#     st.markdown("### Extracted Markdown Preview")
+#     st.write(df.head())
+
+#     # ---------------------------------------------------------------------------------------
+#     # User Input for Topics
+#     # ---------------------------------------------------------------------------------------
+#     st.markdown("### Enter Topics and Descriptions")
+#     num_topics = st.number_input("Number of topics", min_value=1, max_value=10, value=1, step=1)
+
+#     topics = {}
+#     for i in range(num_topics):
+#         topic = st.text_input(f"Topic {i+1} Name", key=f"topic_{i}")
+#         description = st.text_area(f"Topic {i+1} Description", key=f"description_{i}")
+#         if topic and description:
+#             topics[topic] = description
+
+#     # Add a button to execute the analysis
+#     if st.button("Run Analysis"):
+#         if not topics:
+#             st.warning("Please enter at least one topic and description.")
+#             st.stop()
+
+#         # ---------------------------------------------------------------------------------------
+#         # Your existing SurveyAnalysis and extract_excerpts functions remain unchanged here:
+#         # ---------------------------------------------------------------------------------------
+#         analyzer = SurveyAnalysis()
+#         processed_df = analyzer.process_dataframe(df, topics)
+#         df_VIP_extracted = extract_excerpts(processed_df)
+
+#         required_columns = ['Document_Text', 'Topic_Summary', 'Excerpt', 'Topic']
+#         missing_columns = [col for col in required_columns if col not in df_VIP_extracted.columns]
+
+#         if missing_columns:
+#             st.error(f"Missing columns after processing: {missing_columns}")
+#             st.stop()
+
+#         df_VIP_extracted = df_VIP_extracted[required_columns]
+
+#         st.markdown("### Processed Meeting Notes")
+#         st.dataframe(df_VIP_extracted)
+
+#         st.write(f"**Number of meeting notes analyzed:** {len(df)}")
+#         st.write(f"**Number of excerpts extracted:** {len(df_VIP_extracted)}")
+
+#         # CSV download
+#         csv = df_VIP_extracted.to_csv(index=False)
+#         st.download_button(
+#             "Download data as CSV",
+#             data=csv,
+#             file_name='extracted_meeting_notes.csv',
+#             mime='text/csv'
+#         )
+
+#         # Topic distribution visualization
+#         topic_counts = df_VIP_extracted['Topic'].value_counts()
+#         frequency_table = pd.DataFrame({'Topic': topic_counts.index, 'Count': topic_counts.values})
+#         frequency_table['Percentage'] = (frequency_table['Count'] / frequency_table['Count'].sum() * 100).round(0)
+
+#         st.markdown("### Topic Distribution")
+#         st.dataframe(frequency_table)
+
+#         fig, ax = plt.subplots(figsize=(10, 5))
+#         ax.bar(frequency_table['Topic'], frequency_table['Count'], color='#3d9aa1')
+#         ax.set_ylabel('Count')
+#         ax.set_title('Frequency of Topics')
+#         st.pyplot(fig)
+
+# else:
+#     st.info("Please upload a PDF file to begin.") 
+
+
+
 # ---------------------------------------------------------------------------------------
 # Imports and Options
 # ---------------------------------------------------------------------------------------
@@ -9,13 +315,10 @@ import fitz  # PyMuPDF
 import io
 import matplotlib.pyplot as plt
 from PIL import Image
-from mlx_vlm import load, generate
-from mlx_vlm.prompt_utils import apply_chat_template
-from mlx_vlm.utils import load_config, stream_generate
-from docling_core.types.doc.document import DocTagsDocument, DoclingDocument
-
-# Set Streamlit to wide mode
-# st.set_page_config(layout="wide")
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from docling_core.types.doc import DoclingDocument
+from docling_core.types.doc.document import DocTagsDocument
+import torch
 
 # ---------------------------------------------------------------------------------------
 # API Configuration
@@ -168,15 +471,18 @@ with st.sidebar:
     """
     st.markdown(Instructions)
 
-# Load SmolDocling model (mlx_vlm version)
+# Load SmolDocling model using transformers
 @st.cache_resource
 def load_smol_docling():
-    model_path = "ds4sd/SmolDocling-256M-preview"
-    model, processor = load(model_path)
-    config = load_config(model_path)
-    return model, processor, config
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
+    model = AutoModelForVision2Seq.from_pretrained(
+        "ds4sd/SmolDocling-256M-preview",
+        torch_dtype=torch.float32
+    ).to(device)
+    return model, processor
 
-model, processor, config = load_smol_docling()
+model, processor = load_smol_docling()
 
 # Convert PDF to images
 def convert_pdf_to_images(pdf_file):
@@ -190,22 +496,41 @@ def convert_pdf_to_images(pdf_file):
         images.append(image)
     return images
 
-# Extract structured markdown text using SmolDocling (mlx_vlm)
+# Extract structured markdown text using SmolDocling (transformers)
 def extract_markdown_from_image(image):
-    prompt = "Convert this page to docling."
-    formatted_prompt = apply_chat_template(processor, config, prompt, num_images=1)
-    output = ""
-
-    for token in stream_generate(
-        model, processor, formatted_prompt, [image], max_tokens=4096, verbose=False):
-        output += token.text
-        if "</doctag>" in token.text:
-            break
-
-    # Convert DocTags to Markdown
-    doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([output], [image])
+    prompt_text = "Convert this page to docling."
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Prepare inputs
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": prompt_text}
+            ]
+        }
+    ]
+    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=[image], return_tensors="pt").to(device)
+
+    # Generate outputs
+    generated_ids = model.generate(**inputs, max_new_tokens=1024)
+    prompt_length = inputs.input_ids.shape[1]
+    trimmed_generated_ids = generated_ids[:, prompt_length:]
+    doctags = processor.batch_decode(trimmed_generated_ids, skip_special_tokens=False)[0].lstrip()
+
+    # Clean the output
+    doctags = doctags.replace("<end_of_utterance>", "").strip()
+
+    # Populate document
+    doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image])
+
+    # Create a docling document
     doc = DoclingDocument(name="ExtractedDocument")
     doc.load_from_doctags(doctags_doc)
+
+    # Export as markdown
     markdown_text = doc.export_to_markdown()
     return markdown_text
 
@@ -300,4 +625,4 @@ if uploaded_file:
         st.pyplot(fig)
 
 else:
-    st.info("Please upload a PDF file to begin.") 
+    st.info("Please upload a PDF file to begin.")