From 3aaf4f7a9be3c70540eeea254cce380f10f99aa1 Mon Sep 17 00:00:00 2001
From: Jesus Fileto <jesus_filetojr@icloud.com>
Date: Tue, 26 Nov 2024 10:16:36 -0700
Subject: [PATCH] Prevent Cold HuggingFace Models from giving testing errors

---
 libs/chonky/src/embedding.rs | 12 +++++++++++-
 libs/chonky/src/lib.rs       |  2 +-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/libs/chonky/src/embedding.rs b/libs/chonky/src/embedding.rs
index f85e6fad63c..fa1bbc61d80 100644
--- a/libs/chonky/src/embedding.rs
+++ b/libs/chonky/src/embedding.rs
@@ -387,7 +387,12 @@ pub mod hugging_face_api {
     /// [`ChonkyError::HuggingFaceAPI`] when there are HTTP request errors
     pub fn make_table_recognition_request(
         image_path: String,
+        retry: bool,
     ) -> Result<Vec<TablePrediction>, Report<ChonkyError>> {
+        // this is where we would wish to provide add the retry mechanism
+
+        // error code when model is warm is a 503 error, we can then add x-wait-for-model:true for
+        // it to work
         let url =
             "https://api-inference.huggingface.co/models/microsoft/table-transformer-detection";
 
@@ -404,6 +409,11 @@ pub mod hugging_face_api {
             .append(&format!("Authorization: Bearer {access_token}"))
             .change_context(ChonkyError::HuggingFaceAPI)?;
 
+        // we add wait for model to be true if receiving api error prev
+        headers
+            .append(&format!("x-wait-for-model:{retry}"))
+            .change_context(ChonkyError::HuggingFaceAPI)?;
+
         easy.http_headers(headers)
             .change_context(ChonkyError::HuggingFaceAPI)?;
 
@@ -439,7 +449,7 @@ pub mod hugging_face_api {
         fn table_recognition() -> Result<(), Report<ChonkyError>> {
             let file_path = "tests/docs/table-testing.png";
 
-            let table_predictions = make_table_recognition_request(file_path.to_owned())?;
+            let table_predictions = make_table_recognition_request(file_path.to_owned(), true)?;
 
             assert_snapshot!(
                 "table_bounding_boxes.txt",
diff --git a/libs/chonky/src/lib.rs b/libs/chonky/src/lib.rs
index 3f91ad3c130..44625b6a49b 100644
--- a/libs/chonky/src/lib.rs
+++ b/libs/chonky/src/lib.rs
@@ -170,7 +170,7 @@ pub mod pdf_segmentation {
     ) -> Result<Vec<Vec<ExtractedTable>>, Report<ChonkyError>> {
         let mut pdf_table_bounds = Vec::new();
         for (index, page) in pdf.pages().iter().enumerate() {
-            let table_predictions = make_table_recognition_request(images[index].clone())?;
+            let table_predictions = make_table_recognition_request(images[index].clone(), true)?;
 
             let mut page_table_bounds: Vec<ExtractedTable> = Vec::new();
             //convert the pixels back to pdf points