Skip to content

README.md

Latest commit

 

History

History
101 lines (61 loc) · 7.71 KB

README.md

File metadata and controls

101 lines (61 loc) · 7.71 KB
Raw

Persona-Driven Document Intelligence (Adobe Hackathon - Challenge 1B)

This repository contains the source code for a persona-driven document intelligence system developed for the Adobe Hackathon, Challenge 1B. The solution analyzes a collection of PDF documents and acts as an intelligent research assistant, extracting and prioritizing the most relevant sections based on a specific user persona and their "job-to-be-done."

The system is built on a sophisticated hybrid architecture that combines rule-based feature extraction, a lightweight machine learning model for document structuring, and a Small Language Model (SLM) for final content refinement. This approach ensures high accuracy and relevance while adhering to the strict operational constraints of the challenge.

1. Project Objective

The mission of Challenge 1B was to build a generic system that could:

  1. Accept a collection of 3-10 related PDF documents.
  2. Take a persona (e.g., "Investment Analyst") and a job-to-be-done (e.g., "Analyze revenue trends") as input.
  3. Intelligently extract and rank the most relevant sections from the documents that align with the user's goal.
  4. Operate entirely offline, on a CPU-only environment, with a model size limit of 1GB and a processing time limit of 60 seconds for a document collection.

2. Our Development Strategy

Our strategy evolved from initial research to a final, robust hybrid model that balances modern AI capabilities with high efficiency.

Initial Research: The Pure SLM Approach (Scrapped)

We first explored using a Small Language Model (SLM), tinyllama-1.1b-chat, for the entire pipeline. The initial idea was to prompt the SLM to read raw text from each page and directly generate a structured list of headings.

However, this approach was quickly abandoned due to several critical limitations:

  • Performance: The sequential nature of prompting an SLM for every page proved too slow for the 60-second time limit on a CPU.
  • Reliability: The SLM often "hallucinated" headings that weren't in the text or failed to produce perfectly valid JSON, making the output inconsistent and requiring complex error handling.
  • Context Loss: An SLM with a small context window (n_ctx=2048) struggled to understand the overall document structure.

The Final Strategy: A Hybrid AI Pipeline

Recognizing the flaws of the pure SLM approach, we pivoted to a more robust, multi-stage hybrid pipeline. Our hypothesis was that we could achieve superior speed and accuracy by using the right tool for each task:

  1. For Speed & Structure: Use traditional feature engineering and a lightweight ML model.
  2. For Relevance: Use high-performance sentence embeddings.
  3. For Refinement: Use an SLM only for the final, high-value task of summarizing the most relevant content.

This hybrid model proved to be significantly faster, more accurate, and more reliable.

3. Technical Deep Dive: The Hybrid Pipeline

Our final solution is organized into a main script (main_final.py) that executes a two-stage pipeline, supported by helper functions in utils.py.

Stage 1: High-Accuracy Document Structuring

This stage is the foundation of our solution. Instead of relying on an SLM, we built a highly accurate document structuring module to identify all potential headings and their content.

  1. Rich Feature Extraction (extract_text_blocks): We use PyMuPDF for its exceptional speed to parse each PDF and extract not just text, but a rich set of features for every text block, including font size, font name, and a bold flag.

  2. Heading Classification (LightGBM): We use a pre-trained LightGBM classifier to determine if a text block is a heading (H1, H2, H3) or a paragraph. This model was trained on a mix of stylistic and semantic features, making it far more accurate and thousands of times faster than prompting an SLM.

  3. Section Aggregation: The script then iterates through the classified blocks, grouping the paragraph text under its parent heading to create a clean, structured list of all sections in the document collection.

Stage 2: Persona-Based Ranking and Refinement

Once the documents are structured, the second stage focuses on finding and refining the most relevant content for the user.

  1. Semantic Ranking (ranking): We use the SentenceTransformer library with the all-MiniLM-L6-v2 model. This model, known for its excellent balance of performance and speed, converts the user's query ("As a {persona}, I need to {jtbd}") and each extracted section into vector embeddings. By calculating the cosine similarity between the query and each section, we can rank all sections in the document collection by their relevance to the user's goal.

  2. Content Refinement (refine_sections): For the top-ranked sections, we leverage the tinyllama-1.1b-chat SLM in a targeted manner. A carefully crafted prompt asks the SLM to act as the specified persona and provide a concise, one-sentence summary of the key insight from the section's text. This gives the user a quick, actionable summary without the performance overhead of using the SLM on the entire document.

  3. Final JSON Formatting: The final, ranked, and refined data is structured into the exact JSON format required by the challenge submission guidelines.

4. Tech Stack & Justification

Every tool was chosen to maximize performance while respecting the hackathon's constraints.

  • PyMuPDF: Chosen for its raw speed and ability to extract low-level block metadata (fonts, styles), which is essential for our feature-based heading classifier.
  • LightGBM: The ideal classifier for this task. Its high performance and low memory usage on CPUs were critical for meeting the time and resource limits.
  • Sentence-Transformers: Provides state-of-the-art semantic understanding in a compact, offline package that runs exceptionally fast on a CPU.
  • Llama-CPP-Python: Allows us to run a GGUF-quantized version of TinyLLAMA efficiently on a CPU for the final refinement step, keeping us well within the 1GB model size limit.
  • Scikit-Learn / Joblib: Used for creating the ML pipeline and for serializing the trained LightGBM model and other processing objects.
  • Docker: Used to create a self-contained, reproducible environment that precisely matches the competition's evaluation setup.

5. How to Build and Run

Our solution is fully containerized and is designed to be run using Docker.

Build Command

The Docker image is built for the linux/amd64 architecture with the following command:

docker build --platform linux/amd64 -t persona-document-intelligence:latest .

Run Command

Execute the container with the following command. The container will automatically process an input JSON and all associated PDFs from an ./input directory and save the results to an ./output directory.

docker run --rm -v $(pwd)/input:/app/input -v $(pwd)/output:/app/output --network none persona-document-intelligence:latest

6. Code Structure

  • main_final.py: The main entry point of the application. It orchestrates the entire pipeline from loading data to extracting, ranking, refining, and saving the final output.
  • utils.py: Contains helper functions for text cleaning, feature extraction, and formatting the final JSON output.
  • Dockerfile: Defines the instructions to build the Docker image, including setting up the environment and installing all dependencies from requirements.txt.
  • requirements.txt: A list of all Python libraries required to run the project.
  • /models: A directory containing the serialized LightGBM model, the SentenceTransformer model, and the tinyllama GGUF model file.