added json invoices

app.py

@@ -32,6 +32,7 @@ st.set_page_config(page_title="Invoice generator", layout="wide")
 output_folder = "output"
 data_folder = "data"
 template = "template.md"
+invoice_json_path = f"{data_folder}/invoices.json"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 with open(f"{data_folder}/car_parts.json", "r") as f:
     car_parts = json.load(f)
@@ -51,6 +52,7 @@ if "pipeline" not in st.session_state:
         path_to_template=f"{data_folder}/{template}",
         reranker_model="monovlm",
         device=device,
+        invoice_json_path=invoice_json_path,
         gpu_memory_utilization=0.65
     )
 pipeline = st.session_state.pipeline

src/RAG.py

@@ -10,6 +10,7 @@ import torch
 import numpy as np
 from tqdm import tqdm
 import base64
+import json
 
 
 class RAG:
@@ -21,7 +22,8 @@ class RAG:
             device="cpu",
             image_invoice_index_path=None,
             path_to_invoices=None,
-            path_to_images=None
+            path_to_images=None,
+            path_to_invoice_json=None
     ):
         self.index = faiss.read_index(fais_index_path)
         self.model, self.preprocess = clip.load(clip_model, device=device)
@@ -31,6 +33,9 @@ class RAG:
         self.path_to_invoices = path_to_invoices
         self.path_to_images = path_to_images
         self.reranker = reranker
+        if path_to_invoice_json:
+            with open(path_to_invoice_json, "r") as f:
+                self.invoice_json = json.load(f)
 
     @staticmethod
     def image_to_base64(image_path):
@@ -56,7 +61,45 @@ class RAG:
         distances, indices = self.index.search(image_features, k)
         return distances, indices
 
-    def find_invoice(self, image=None, image_path=None, return_only_path=True, k=1, damage_description=None):
+    def return_invoice_table(self, path=None, invoice_is_table=True):
+        if path is None and not invoice_is_table:
+            raise ValueError("Path to invoice must be provided.")
+        if self.invoice_json is None and invoice_is_table:
+            raise ValueError("Path to invoice json must be provided.")
+
+        if invoice_is_table:
+            return self.invoice_json[path]
+
+        pdf_path = f"{self.path_to_invoices}/{path}"
+        reader = PdfReader(pdf_path)
+        page = reader.pages[0]
+        text = page.extract_text()
+
+        table_text = re.search(r"Beschädigtes Teil.*?Gesamtsumme:.*?EUR", text, re.DOTALL).group()
+
+        lines = table_text.splitlines()
+        header = lines[0]
+        other_text = "\n".join(lines[1:])
+        cleaned_text = re.sub(r"(?<!\d)\n", " ", other_text)
+
+        table = header + "\n" + cleaned_text
+
+        inv = table.split("\n")
+        reformatted_inv = "Beschädigtes Teil | Teilkosten (EUR) | Arbeitsstunden | Arbeitskosten (EUR/Stunde) | Gesamtkosten (EUR)\n" + "\n".join(
+            " ".join(inv[i].split(" ")[:-4]) + " | " + ' | '.join(inv[i].split(" ")[-4:]) for i in
+            range(1, len(inv) - 1)) + "\n" + inv[-1]
+
+        return reformatted_inv
+
+    def find_invoice(
+            self,
+            image=None,
+            image_path=None,
+            return_only_path=True,
+            k=1,
+            damage_description=None,
+            invoice_is_table=True
+    ):
         if self.image_invoice_index is None:
             raise ValueError("No index for invoices found.")
         _, indices = self.search_image(image=image, image_path=image_path, k=k)
@@ -67,8 +110,6 @@ class RAG:
         if self.reranker:
             if damage_description is None:
                 raise ValueError("Damage description must be provided.")
-            # images = [self.image_to_base64(f"{self.path_to_images}/{img_path}") for img_path in images_paths]
-
             images = [f"{self.path_to_images}/{img_path}" for img_path in images_paths]
             results = self.reranker.rank(damage_description, images, doc_ids=invoices)
             invoices = [doc.doc_id for doc in results]
@@ -83,25 +124,7 @@ class RAG:
         invoices_tables = []
 
         for invoice in invoices:
-            pdf_path = f"{self.path_to_invoices}/{invoice}"
-            reader = PdfReader(pdf_path)
-            page = reader.pages[0]
-            text = page.extract_text()
-
-            table_text = re.search(r"Beschädigtes Teil.*?Gesamtsumme:.*?EUR", text, re.DOTALL).group()
-
-            lines = table_text.splitlines()
-            header = lines[0]
-            other_text = "\n".join(lines[1:])
-            cleaned_text = re.sub(r"(?<!\d)\n", " ", other_text)
-
-            table = header + "\n" + cleaned_text
-
-            inv = table.split("\n")
-            reformatted_inv = "Beschädigtes Teil | Teilkosten (EUR) | Arbeitsstunden | Arbeitskosten (EUR/Stunde) | Gesamtkosten (EUR)\n" + "\n".join(
-                " ".join(inv[i].split(" ")[:-4]) + " | " + ' | '.join(inv[i].split(" ")[-4:]) for i in
-                range(1, len(inv) - 1)) + "\n" + inv[-1]
-
+            reformatted_inv = self.return_invoice_table(invoice, invoice_is_table)
             invoices_tables.append(reformatted_inv)
 
         return invoices_tables, invoices

src/pipelines.py

@@ -18,6 +18,7 @@ class InvoiceGenerator:
             path_to_template,
             reranker_model=None,
             device="cuda",
+            invoice_json_path=None,
             max_model_len=4096, max_tokens=2048, gpu_memory_utilization=0.95
     ):
         self.model = Pixtral(max_model_len=max_model_len, max_tokens=max_tokens,
@@ -31,7 +32,8 @@ class InvoiceGenerator:
             image_invoice_index_path=image_invoice_index_path,
             path_to_invoices=path_to_invoices,
             path_to_images=path_to_images,
-            reranker=self.reranker
+            reranker=self.reranker,
+            path_to_invoice_json=invoice_json_path
         )
         self.path_to_invoices = path_to_invoices
         self.path_to_images = path_to_images