Update app.py

app.py

@@ -1,52 +1,27 @@
 import PyPDF2
 from pprint import pprint
+from getpass import getpass
 from haystack import Pipeline
 from haystack.schema import Document
 from haystack.nodes import BM25Retriever
 from haystack.document_stores import InMemoryDocumentStore
-from haystack.nodes import PreProcessor, PromptTemplate, PromptNode
-from pdf2image import convert_from_path
-import pytesseract
-from PIL import Image
+from haystack.nodes import PromptTemplate, PromptNode
 import gradio as gr
 import os
-from pydantic import BaseModel
 
-# Function to extract text from a PDF file using OCR
-def extract_text_from_pdf(pdf_path):
-    text = ""
-    # Convert PDF pages to images
-    images = convert_from_path(pdf_path)
-    for image in images:
-        # Perform OCR on the image
-        text += pytesseract.image_to_string(image)
-    return text
-
-class Config(BaseModel):
-    class Config:
-        arbitrary_types_allowed = True
+HF_TOKEN = getpass("Enter Token")
+from huggingface_hub import notebook_login
+notebook_login()
 
 # Process and retrieve answers
-def process_invoice(pdf, hf_token, questions):
-    # Extract text from the PDF
-    extracted_text = extract_text_from_pdf(pdf.name)
-    document = Document(content=extracted_text)
+def process_invoice(file, hf_token, questions):
+    # Read file content
+    file_content = file.read()
+    document = Document(content=file_content)
     docs = [document]
 
-    # Initializing the processor
-    processor = PreProcessor(
-        clean_empty_lines=True,
-        clean_whitespace=True,
-        clean_header_footer=True,
-        split_by="word",
-        split_length=500,
-        split_respect_sentence_boundary=True,
-        split_overlap=0,
-    )
-
-    preprocessed_docs = processor.process(docs)
     document_store = InMemoryDocumentStore(use_bm25=True)
-    document_store.write_documents(preprocessed_docs)
+    document_store.write_documents(docs)
     retriever = BM25Retriever(document_store, top_k=2)
 
     qa_template = PromptTemplate(prompt=
@@ -78,20 +53,20 @@ def process_invoice(pdf, hf_token, questions):
     return answers
 
 # Gradio interface
-def gradio_interface(pdf, hf_token, questions):
-    answers = process_invoice(pdf, hf_token, questions)
+def gradio_interface(file, hf_token, questions):
+    answers = process_invoice(file, hf_token, questions)
     return answers
 
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=[
-        gr.inputs.File(file_count="single", type="file", label="Upload Invoice (PDF)"),
+        gr.inputs.File(file_count="single", type="file", label="Upload Invoice (PDF or Image)"),
         gr.inputs.Textbox(type="password", label="Enter your Hugging Face Token"),
         gr.inputs.Textbox(lines=5, placeholder="Enter your questions separated by commas")
     ],
     outputs="json",
     title="Invoice Data Extraction",
-    description="Upload an invoice PDF, provide your Hugging Face token, and get the extracted data based on your questions."
+    description="Upload an invoice PDF or image, provide your Hugging Face token, and get the extracted data based on your questions."
 )
 
 if __name__ == "__main__":