Saturday, October 12, 2024

From Concept to PRD: My Journey Collaborating with AI on a RAG-based Extraction System

I recently embarked on an exciting project to overhaul our document attribute extraction system. What started as a simple idea quickly evolved into a comprehensive plan for a cutting-edge RAG-based system. Throughout this journey, I collaborated with an AI assistant, and I want to share how this partnership helped shape our project roadmap.

The Initial Concept


It all began when I approached our AI assistant with a straightforward question: "Code to fine tune a llama 3.2 model for nested structured data extraction from pdf files. What should be the training dataset format?"

Little did I know that this simple query would spark a series of discussions that would completely transform our approach to document extraction.

Embracing RAG: A Game-Changer

As we delved deeper into the possibilities, the AI suggested implementing a Retrieval-Augmented Generation (RAG) approach. This concept immediately piqued my interest. We explored how RAG could enhance our system's ability to handle complex, lengthy documents while maintaining high accuracy.

The AI provided a detailed explanation of how we could structure our system:Chunk the input document
  1. Create embeddings using an E5 model
  2. Generate synthetic answers with a fine-tuned LLaMA model
  3. Retrieve and re-rank relevant chunks using ColBERT v2
  4. Extract attributes from the top-ranked chunks

This approach seemed promising, but I had concerns about performance, especially for transient documents that require quick processing.

Optimizing for Speed and Accuracy

Addressing my concerns, we brainstormed ways to optimize the system for transient documents. The AI suggested implementing a "fast track" pipeline that uses lighter models and skips some computationally expensive steps. This solution struck a balance between speed and accuracy, potentially processing transient documents 50-70% faster than the full pipeline.

Expanding Capabilities: Dependent Data Extraction

As we refined our plan, I realized we needed to handle more complex scenarios. I asked, "Can this be used for doing dependent data extraction? Like find a set of ids and extract specific set of attributes for each of those ids?"

The AI's response was enthusiastic and detailed. We worked together to design a two-stage extraction process:
  1. ID Extraction: Identify and extract a set of IDs from the document
  2. Attribute Extraction per ID: Perform targeted attribute extraction for each ID
This feature significantly expanded the versatility of our system, allowing it to handle nested data structures common in many business documents.

Bringing It All Together: The PRD

As our ideas coalesced, I asked the AI to generate a Product Requirements Document (PRD). The resulting document was comprehensive, covering everything from key features and technical requirements to performance metrics and potential risks.

What impressed me most was how the PRD evolved through our conversation. When I requested updates to include new features or address specific concerns, the AI quickly incorporated these changes, resulting in a well-rounded, thoughtful project plan.

Lessons Learned

Reflecting on this experience, I've gained valuable insights into collaborating with AI:

  1. Iterative Refinement: Our initial idea evolved significantly through back-and-forth discussion. Don't be afraid to explore tangents or challenge the AI's suggestions.
  2. Leverage AI's Knowledge: The AI brought up concepts and technologies I hadn't considered, like using E5 for embeddings and ColBERT v2 for re-ranking. This broadened our solution space.
  3. Human Expertise is Crucial: While the AI provided extensive technical knowledge, my understanding of our specific needs and constraints was vital in shaping a practical solution.
  4. AI as a Brainstorming Partner: The AI excelled at generating ideas and fleshing out details, making it an excellent brainstorming partner.
  5. Clarity in Communication: Being clear and specific in my queries led to more targeted and useful responses.

Looking Ahead

This collaboration has set us on an exciting path. We now have a solid plan for a RAG-based document attribute extraction system that promises to be more accurate, flexible, and efficient than our current solution.

As we move into the implementation phase, I'm confident that the groundwork we've laid through this AI-assisted planning process will prove invaluable. It's a testament to how AI can augment human creativity and expertise, leading to more innovative and comprehensive solutions.

The journey from a simple question about dataset formats to a full-fledged PRD for a cutting-edge system has been enlightening. It's clear that AI assistants like the one I worked with are not just tools for answering questions, but partners in the creative and strategic thinking process.

I'm excited to see how this project unfolds and look forward to sharing more insights as we bring our RAG-based extraction system to life!


---


PS: This post and the entire communication was done with Claude Sonnet 3.5 Model.