Part 3: Desktop App for Document QA with RAG - Document Layout Analysis, Text Extraction and Generation

10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

// Imports omitted

// Copied from: https://github.com/styrowolf/layoutparser-ort/blob/master/src/utils.rs
/// Utility function to convert bbox to a array
fn vec_to_bbox<T: Copy>(v: Vec<T>) -> [T; 4] {
    [v[0], v[1], v[2], v[3]]
}

// An emum to represent the classes of regions of interest
// detected by the `layout detection` model
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub enum DetectedElem {
    Text,
    Title,
    List,
    Table,
    Figure,
}

impl Display for DetectedElem {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}",
            match self {
                Self::Text => "Text",
                Self::Title => "Title",
                Self::List => "List",
                Self::Table => "Table",
                Self::Figure => "Figure",
            }
        )
    }
}

/// This struct represents a Region of interest
#[derive(Debug)]
pub struct RegionOfInterest {
    kind: DetectedElem,
    // the bounding box - x1, y1, x2, y2 - top, left, bottom, right
    bbox: [f32; 4],
    // confidence: f32,
}

impl RegionOfInterest {
    pub fn kind(&self) -> DetectedElem {
        self.kind
    }

    pub fn bbox(&self) -> [f32; 4] {
        self.bbox
    }
}

 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208

// Code omitted ..

/// A [`Detectron2`](https://github.com/facebookresearch/detectron2)-based model.
pub struct Detectron2Model {
    model: Session,
    label_map: [DetectedElem; 5],
}

// Copied from: https://github.com/styrowolf/layoutparser-ort/blob/master/src/utils.rs
/// Utility function to convert bbox to a array
fn vec_to_bbox<T: Copy>(v: Vec<T>) -> [T; 4] {
    [v[0], v[1], v[2], v[3]]
}

impl Detectron2Model {
    /// Required input image width.
    pub const REQUIRED_WIDTH: usize = 800;
    /// Required input image height.
    pub const REQUIRED_HEIGHT: usize = 1035;
    /// Default confidence threshold for detections.
    pub const DEFAULT_CONFIDENCE_THRESHOLD: f32 = 0.85;

    pub fn new() -> Result<Self> {
        // Loading and initializing the model from `onnx` file
        let model = Session::builder()?
            .with_optimization_level(GraphOptimizationLevel::Level3)?
            // We could make this a little more generic with `numcpus` crate
            .with_intra_threads(8)?
            .commit_from_file("../models/layout.onnx")?;

        // You could print the model outputs to figure out which prediction datapoints are useful
        // println!("{:?}", model.outputs);

        Ok(Self {
            model,
            label_map: [
                DetectedElem::Text,
                DetectedElem::Title,
                DetectedElem::List,
                DetectedElem::Table,
                DetectedElem::Figure,
            ],
        })
    }

    pub fn predict(&self, page: &image::DynamicImage) -> Result<Vec<RegionOfInterest>> {
        let (img_width, img_height, input) = self.preprocess(page)?;
        // let hm = HashMap::from([("x.1".to_string(), input)]);
        let res = self.model.run(ort::inputs!["x.1" => input]?)?;

        self.postprocess(res, img_width, img_height)
    }

    // 1. Resizes an image to the required format!
    // 2. Creates a tensor from the image
    // 3. Reshapes the tensor to channel first format
    // 4. Creates input ndarray for `ort` to consume
    fn preprocess(&self, img: &image::DynamicImage) -> Result<(u32, u32, ort::Value)> {
        // TODO: re-visit this and resize smarter
        let (img_width, img_height) = (img.width(), img.height());
        let img = img.resize_exact(
            Self::REQUIRED_WIDTH as u32,
            Self::REQUIRED_HEIGHT as u32,
            imageops::FilterType::Triangle,
        );

        let img = img.to_rgb8().into_raw();

        // Read the image as a tensor
        let t = Tensor::from_vec(
            img,
            (Self::REQUIRED_HEIGHT, Self::REQUIRED_WIDTH, 3),
            &Device::Cpu,
        )?
        .to_dtype(DType::F32)?
        .permute((2, 0, 1))? // shape: [3, height, width]
        .to_vec3::<f32>()?
        .concat()
        .concat();

        // Create a `ndarray` input for `ort` runtime to consume
        let input = ort::Value::from_array(
            ([3, Self::REQUIRED_HEIGHT, Self::REQUIRED_WIDTH], &t[..])
        )?;
        
        Ok((img_width, img_height, input.into()))
    }

    // Reads the predictions and converts them to regions of interest
    fn postprocess(
        &self,
        outputs: SessionOutputs<'_, '_>,
        width: u32,
        height: u32,
    ) -> Result<Vec<RegionOfInterest>> {
        // Extract predictions for bounding boxes,
        // labels and confidence scores
        // Shape: [num pred, 4]
        let bboxes = &outputs[0].try_extract_tensor::<f32>()?;
        // Shape: [num pred]
        let labels = &outputs[1].try_extract_tensor::<i64>()?;
         // 3 for MASK_RCNN_X_101_32X8D_FPN_3x | 2 for FASTER_RCNN_R_50_FPN_3X
         // Shape: [num pred]
        let confidence = &outputs[3].try_extract_tensor::<f32>()?;

        // We had originally `resized` the image to fit
        // the required input dimensions,
        // we are just going to adjust the predictions to factor in the resize
        let width_factor = width as f32 / Self::REQUIRED_WIDTH as f32;
        let height_factor = height as f32 / Self::REQUIRED_HEIGHT as f32;

        // Iterate over (region bounding boxes, predicted classes/ labels, and confidence scores)
        let mut elements = bboxes
            .rows()
            .into_iter()
            .zip(labels.iter().zip(confidence.iter()))
            .filter_map(|(bbox, (&label, &confidence))| {
                // Skip everything below some confidence score we want to work with
                if confidence < Self::DEFAULT_CONFIDENCE_THRESHOLD {
                    return None;
                }

                // Getting the predicted label from the predicted index
                let label = self.label_map.get(label as usize)?;
                // We don't have any way of interpreting Figure and Table as text
                // So, we'll skip that
                if label == &DetectedElem::Figure || label == &DetectedElem::Table {
                    return None;
                }
                let [x1, y1, x2, y2] = vec_to_bbox(bbox.iter().copied().collect::<Vec<_>>());
                // Adjusting the predicted bounding box to our original image size
                Some(RegionOfInterest {
                    kind: *label,
                    bbox: [
                        x1 * width_factor,
                        y1 * height_factor,
                        x2 * width_factor,
                        y2 * height_factor,
                    ],
                    confidence
                })
            })
            .collect::<Vec<_>>();

        // Now we sort the predictions to (kind of) visual hierarchy
        // from top left
        elements.par_sort_unstable_by(|a, b| {
            (a.bbox()[1].max(a.bbox()[3])).total_cmp(&(b.bbox()[1].max(b.bbox()[3])))
        });

        Ok(elements)
    }
}

214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230

#[cfg(test)]
mod tests {
    use anyhow::Result;

    use super::Detectron2Model;

    #[test]
    fn single_page_layout() -> Result<()> {
        let d2model = Detectron2Model::new()?;
        let img = image::open("../test-data/paper-image.jpg")?;

        let pred = d2model.predict(&img)?;
        println!("{pred:?}");

        Ok(())
    }
}

28
29
30
31
32
33
34
35

// Imports and code omitted

pub struct PdfProc {
    pdfs: Vec<PathBuf>,
    layout: Detectron2Model,
    pdfium: Pdfium,
    pdfium_cfg: PdfRenderConfig,
}

 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

// code omitted ..
impl PdfProc {
    // initializer ..
    pub fn new(model_path: &Path, pdfs: Vec<PathBuf>) -> Result<Self> {
        // Straitforward implementation for initializing the different fields of the struct
        Ok(Self { ... })
    }

    /// Returns the total number of pages to be analyzed
    pub fn estimate(&self) -> usize {
        // this can be used to show some progress in real world
    }

    /// Extract text from `pdfs`
    pub fn extract(&self, send: Sender<ExtractorEvt>) -> Result<Vec<Vec<(String, FileKind)>>> {
        // for each `.pdf` file we are going to convert the pages to images
        let file_encoded = self
            .pdfs
            .iter()
            .filter_map(|file| {
                let pdf = self.pdfium.load_pdf_from_file(&file, None).ok()?;

                self.process_pages(file, pdf, send.clone())
            })
            .collect::<Vec<_>>();

        Ok(file_encoded)
    }

    /// Processes each pages
    /// - reneders page with rendering config
    /// - runs layout detection
    /// - reads text from bounding boxes detected by the model
    pub fn process_pages(
        &self,
        file: &PathBuf,
        doc: PdfDocument<'_>,
        send: Sender<ExtractorEvt>,
    ) -> Option<Vec<(String, FileKind)>> {
        Some(
            doc.pages()
                .iter()
                .enumerate()
                .filter_map(|(idx, page)| {
                    // Convert the page to an image
                    let img = page.render_with_config(&self.pdfium_cfg).ok()?.as_image(); // Renders this page to an image::DynamicImage...

                    // Keep track of the factors by which the page and images of pages were resized to
                    // This is required to get accurate output from the predicted regions of interest
                    let w_f = page.width().value / img.width() as f32;
                    let h_f = page.height().value / img.height() as f32;
                    let pg_num = idx + 1;

                    // send the image for prediction
                    // and for each predicted `region of interest`
                    // fetch the text inside the bounding box
                    let text = self
                        .layout
                        .predict(&img)
                        .ok()?
                        .iter()
                        .filter_map(|e| {
                            // The bounding box for the region of interest
                            let bbox = e.bbox(); // x1, y1, x2, y2

                            // The bounding boxes for the predicted regions follow a `left-top` co-ordinate system
                            // But `pdfium` uses a bottom-left coordinate system, let's convert it
                            // We'll also factor in the original page size here
                            let top = page.height().value - bbox[1] * h_f + PADDING;
                            let bottom = page.height().value - bbox[3] * h_f - PADDING;
                            let left = bbox[0] * w_f - PADDING;
                            let right = bbox[2] * w_f + PADDING;

                            // Now, we have the `pdfium` compatible bounding boxes
                            // Let's fetch the text
                            let text = page
                                .text()
                                .ok()?
                                .inside_rect(PdfRect::new_from_values(bottom, left, top, right))
                                .replace("\t", " ")
                                .replace("\r\n", "\n");

                            Some(match e.kind() {
                                // We are using `MarkDownSplitter` for our text splitting task
                                // Here we are adding `##` to mark the generated text as title
                                DetectedElem::Title => {
                                    format!("## {}\n", text.replace("\n", "; "))
                                }
                                // Rest of the text remains as is
                                DetectedElem::Text | DetectedElem::List => text,
                                _ => unimplemented!(),
                            })
                        })
                        .collect::<Vec<_>>()
                        .join("\n");
                    
                    if let Err(e) = send.send(ExtractorEvt::Page) {
                        eprintln!("Warn: error sending page event: {e:?}");
                    }
                    
                    Some((text, FileKind::Pdf((file.to_owned(), pg_num))))
                })
                .collect::<Vec<_>>(),
        )
    }
}

 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

// Imports omitted ..

// Sampling constants
const TEMPERATURE: f64 = 0.8;
const TOP_P: f64 = 0.95;
const TOP_K: usize = 32;

/// A struct to maintain a initialized Llama quantized `gguf` model and associated methods
pub struct Generator {
    cfg: Config,
    device: Device,
    model: Llama,
    tokenizer: Tokenizer,
    sampler: LogitsProcessor,
    stop_tokens: [u32; 2],
}

impl Generator {
    // Download model `safetensor` files into your project dir `models` folder
    // I'm using LLaMA3.1 8B instruct, you can use whatever you want
    const MODEL_FILES: [&'static str; 2] = [
        "model-00001-of-00002.safetensors",
        "model-00002-of-00002.safetensors",
    ];
    const TOKENIZER_FILE: &'static str = "llama_tokenizer.json";
    const MODEL_CONFIG_FILE: &'static str = "llama_config.json";

    /// Initializer for new llama manager
    pub fn new(dir: &Path, device: &Device) -> Result<Self> {
        let mut device = device.to_owned();
        if let candle_core::Device::Metal(mut m) = device {
            m.set_use_mlx_mm(false);
            device = Device::Metal(m);
        }

        let (model, mut cfg, tokenizer) = Self::load_model(dir, &device)?;
        let stop_tokens = [
            tokenizer.token_to_id("<|eot_id|>").unwrap(),
            tokenizer.token_to_id("<|end_of_text|>").unwrap(),
        ];
        cfg.max_position_embeddings = 4096;

        // Initializing the sampler
        let sampler = LogitsProcessor::from_sampling(
            42,
            Sampling::TopKThenTopP {
                k: TOP_K,
                p: TOP_P,
                temperature: TEMPERATURE,
            },
        );
        
        println!("Llama ready!");
        Ok(Self {
            cfg,
            device: device.clone(),
            model,
            tokenizer,
            sampler,
            // sampler2,
            stop_tokens,
        })
    }

    // A utility function to load the model and tokenizer
    fn load_model(model_dir: &Path, device: &Device) -> Result<(Llama, Config, Tokenizer)> {
        // let model_file = model_dir.join(Self::MODEL_FILE);
        let tok_file = model_dir.join(Self::TOKENIZER_FILE);
        let cfg_file = model_dir.join(Self::MODEL_CONFIG_FILE);
        let model_files = Self::MODEL_FILES
            .iter()
            .map(|mf| model_dir.join(mf))
            .collect::<Vec<_>>();

        println!("Loading LLaMA ..");
        let start = Instant::now();
        let cfg =
            serde_json::from_slice::<LlamaConfig>(&std::fs::read(&cfg_file)?)?.into_config(false);
        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_files, DType::BF16, device)? };
        let llama = Llama::load(vb, &cfg)?;
        println!("LLaMA loaded in {}s", (Instant::now() - start).as_secs());

        let tokenizer = Tokenizer::from_file(tok_file).unwrap();

        Ok((llama, cfg, tokenizer))
    }

    // Utility function to run the generation loop
    fn generate(&mut self, prompt: &str) -> Result<String> {
        // Tokenize the input
        let input = self
            .tokenizer
            .encode(prompt, true)
            .map_err(|e| anyhow!(e))?;

        if input.len() >= self.cfg.max_position_embeddings {
            return Err(anyhow!("large input tokens!"));
        }

        let mut cache = Cache::new(true, DType::BF16, &self.cfg, &self.device)?;
        // Creating a tensor of input tokens
        let mut ip = Tensor::new(input.get_ids(), &self.device)?.unsqueeze(0)?;

        let mut start = std::time::Instant::now();

        // The forward pass to the first token
        let mut logits = self.model.forward(&ip, 0, &mut cache)?;
        // Sampling the first token
        let mut next = self.sampler.sample(&logits.squeeze(0)?)?;
        
        println!(
            "{} prompt tokens processed @ {}t/s",
            input.len(),
            input.len() as f32 / (std::time::Instant::now() - start).as_secs() as f32
        );
        // A container for all tokens generated
        let mut all_tokens = vec![next];

        start = std::time::Instant::now();

        // Forward pass - decoder loop
        for i in input.len()..self.cfg.max_position_embeddings {
            ip = Tensor::new(&[next], &self.device)?.unsqueeze(0)?;

            logits = self.model.forward(&ip, i, &mut cache)?;
            next = self.sampler.sample(&logits.squeeze(0)?).unwrap();
            if self.stop_tokens.contains(&next) {
                break;
            }

            all_tokens.push(next);
        }
        println!(
            "{} tokens generated @ {}t/s",
            all_tokens.len() - 1,
            (all_tokens.len() - 1) as f32 / (std::time::Instant::now() - start).as_secs_f32()
        );

        // Decode tokens and return result
        Ok(match self.tokenizer.decode(&all_tokens[..], false) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Error generating tokens: {e:?}");
                anyhow::bail!("Error generating tokens")
            }
        })
    }
}