# --------------------------------------------------------------------------------------- # 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 transformers import AutoProcessor, AutoModelForVision2Seq from docling_core.types.doc import DoclingDocument from docling_core.types.doc.document import DocTagsDocument import torch import os from huggingface_hub import InferenceClient # --------------------------------------------------------------------------------------- # Streamlit 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" } ) # --------------------------------------------------------------------------------------- # Streamlit Sidebar # --------------------------------------------------------------------------------------- st.sidebar.title("📌 About This App") st.sidebar.markdown(""" #### ⚠️ **Important Note on Processing Time** This app uses the **SmolDocling** model (`ds4sd/SmolDocling-256M-preview`) to convert PDF pages into markdown text. Currently, the model is running on a CPU-based environment (**CPU basic | 2 vCPU - 16 GB RAM**), and therefore processing each page can take a significant amount of time (approximately **6 minutes per page**). **Note: It is recommended that you upload single-page PDFs, as testing showed approximately 6 minutes of processing time per page.** This setup is suitable for testing and demonstration purposes, but **not efficient for real-world usage**. For faster processing, consider running the optimized version `ds4sd/SmolDocling-256M-preview-mlx-bf16` locally on a MacBook, where it performs significantly faster. --- #### 🛠️ **How This App Works** Here's a quick overview of the workflow: 1. **Upload PDF**: You upload a PDF document using the uploader provided. 2. **Convert PDF to Images**: The PDF is converted into individual images (one per page). 3. **Extract Markdown from Images**: Each image is processed by the SmolDocling model to extract markdown-formatted text. 4. **Enter Topics and Descriptions**: You provide specific topics and their descriptions you'd like to extract from the document. 5. **Extract Excerpts**: The app uses the **meta-llama/Llama-3.1-70B-Instruct** model to extract exact quotes relevant to your provided topics. 6. **Results in a DataFrame**: All extracted quotes and their topics are compiled into a structured DataFrame that you can preview and download. --- Please proceed by uploading your PDF file to begin the analysis. """) # --------------------------------------------------------------------------------------- # Session State Initialization # --------------------------------------------------------------------------------------- for key in ['pdf_processed', 'markdown_texts', 'df']: if key not in st.session_state: st.session_state[key] = False if key == 'pdf_processed' else [] # --------------------------------------------------------------------------------------- # API Configuration # --------------------------------------------------------------------------------------- hf_api_key = os.getenv('HF_API_KEY') if not hf_api_key: raise ValueError("HF_API_KEY not set in environment variables") client = InferenceClient(api_key=hf_api_key) # --------------------------------------------------------------------------------------- # Survey Analysis Class # --------------------------------------------------------------------------------------- class AIAnalysis: def __init__(self, client): self.client = client def prepare_llm_input(self, document_content, topics): topic_descriptions = "\n".join([f"- **{t}**: {d}" for t, d in topics.items()]) return f"""Extract and summarize PDF notes based on topics: {topic_descriptions} Instructions: - Extract exact quotes per topic. - Ignore irrelevant topics. - Strictly follow this format: [Topic] - "Exact quote" Document Content: {document_content} """ def prompt_response_from_hf_llm(self, llm_input): system_prompt = """ You are an expert assistant tasked with extracting exact quotes from provided meeting notes based on given topics. Instructions: - Only extract exact quotes relevant to provided topics. - Ignore irrelevant content. - Strictly follow this format: [Topic] - "Exact quote" """ response = self.client.chat.completions.create( model="meta-llama/Llama-3.1-70B-Instruct", messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": llm_input} ], stream=True, temperature=0.5, max_tokens=1024, top_p=0.7 ) response_content = "" for message in response: # Correctly handle streaming response response_content += message.choices[0].delta.content print("Full AI Response:", response_content) # Debugging return response_content.strip() def extract_text(self, response): return response def process_dataframe(self, df, topics): results = [] for _, row in df.iterrows(): llm_input = self.prepare_llm_input(row['Document_Text'], topics) response = self.prompt_response_from_hf_llm(llm_input) notes = self.extract_text(response) results.append({'Document_Text': row['Document_Text'], 'Topic_Summary': notes}) return pd.concat([df.reset_index(drop=True), pd.DataFrame(results)['Topic_Summary']], axis=1) # --------------------------------------------------------------------------------------- # Helper Functions # --------------------------------------------------------------------------------------- @st.cache_resource def load_smol_docling(): 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 = load_smol_docling() def convert_pdf_to_images(pdf_file, dpi=150, max_size=1600): images = [] doc = fitz.open(stream=pdf_file.read(), filetype="pdf") for page in doc: pix = page.get_pixmap(dpi=dpi) img = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB") img.thumbnail((max_size, max_size), Image.LANCZOS) images.append(img) return images def extract_markdown_from_image(image): device = "cuda" if torch.cuda.is_available() else "cpu" prompt = processor.apply_chat_template([{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Convert this page to docling."}]}], add_generation_prompt=True) inputs = processor(text=prompt, images=[image], return_tensors="pt").to(device) with torch.no_grad(): generated_ids = model.generate(**inputs, max_new_tokens=1024) doctags = processor.batch_decode(generated_ids[:, inputs.input_ids.shape[1]:], skip_special_tokens=False)[0].replace("", "").strip() doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image]) doc = DoclingDocument(name="ExtractedDocument") doc.load_from_doctags(doctags_doc) return doc.export_to_markdown() # Revised extract_excerpts function with improved robustness def extract_excerpts(processed_df): rows = [] for _, r in processed_df.iterrows(): sections = re.split(r'\n(?=(?:\*\*|\[)?[A-Za-z/ ]+(?:\*\*|\])?\n- )', r['Topic_Summary']) for sec in sections: topic_match = re.match(r'(?:\*\*|\[)?([A-Za-z/ ]+)(?:\*\*|\])?', sec.strip()) if topic_match: topic = topic_match.group(1).strip() excerpts = re.findall(r'- "?([^"\n]+)"?', sec) for excerpt in excerpts: rows.append({ 'Document_Text': r['Document_Text'], 'Topic_Summary': r['Topic_Summary'], 'Excerpt': excerpt.strip(), 'Topic': topic }) print("Extracted Rows:", rows) # Debugging return pd.DataFrame(rows) # --------------------------------------------------------------------------------------- # 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 and not st.session_state['pdf_processed']: with st.spinner("Processing PDF..."): images = convert_pdf_to_images(uploaded_file) markdown_texts = [extract_markdown_from_image(img) for img in images] st.session_state['df'] = pd.DataFrame({'Document_Text': markdown_texts}) st.session_state['pdf_processed'] = True st.success("PDF processed successfully!") if st.session_state['pdf_processed']: st.markdown("### Extracted Text Preview") st.write(st.session_state['df'].head()) st.markdown("### Enter Topics and Descriptions") num_topics = st.number_input("Number of topics", 1, 10, 1) topics = {} for i in range(num_topics): topic = st.text_input(f"Topic {i+1} Name", key=f"topic_{i}") desc = st.text_area(f"Topic {i+1} Description", key=f"description_{i}") if topic and desc: topics[topic] = desc if st.button("Run Analysis"): if not topics: st.warning("Please enter at least one topic and description.") st.stop() analyzer = AIAnalysis(client) processed_df = analyzer.process_dataframe(st.session_state['df'], topics) extracted_df = extract_excerpts(processed_df) st.markdown("### Extracted Excerpts") st.dataframe(extracted_df) csv = extracted_df.to_csv(index=False) st.download_button("Download CSV", csv, "extracted_notes.csv", "text/csv") if not extracted_df.empty: topic_counts = extracted_df['Topic'].value_counts() fig, ax = plt.subplots() topic_counts.plot.bar(ax=ax, color='#3d9aa1') st.pyplot(fig) else: st.warning("No topics were extracted. Please check the input data and topics.") if not uploaded_file: st.info("Please upload a PDF file to begin.")