Update app.py

app.py

@@ -1,29 +1,35 @@
 import streamlit as st
 from PIL import Image
-from transformers import LayoutLMv3ForTokenClassification, LayoutLMv3Processor
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-
-# Load LayoutLMv3 model and processor
-model_name = "microsoft/layoutlmv3-base"
-layoutlm_model = LayoutLMv3ForTokenClassification.from_pretrained(model_name)
-layoutlm_processor = LayoutLMv3Processor.from_pretrained(model_name)
-
-def extract_text_from_image(image, question):
-    # Preprocess the image using the processor
-    encoding = layoutlm_processor(images=image, return_tensors="pt")
-
-    # Get model predictions
+import easyocr
+
+# Load the question-answering model and tokenizer
+model_name = "t5-base"
+qa_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+qa_tokenizer = AutoTokenizer.from_pretrained(model_name)
+reader = easyocr.Reader(['en'])
+
+# Function to extract text from image using easyocr
+def extract_text_from_image(image):
+    # Perform OCR on the image using easyocr
+    ocr_result = reader.readtext(image, detail=0)
+    text = " ".join(ocr_result)
+    return text
+
+# Function to get response from the language model
+def get_response_from_llm(extracted_text, question):
+    # Prepare the input for the model
+    input_text = f"question: {question} context: {extracted_text}"
+    inputs = qa_tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
+
+    # Generate the response
     with torch.no_grad():
-        outputs = layoutlm_model(**encoding)
-
-    # Get the predicted tokens and labels
-    predicted_tokens = outputs.logits.argmax(-1).squeeze().tolist()
-    tokens = layoutlm_processor.tokenizer.convert_ids_to_tokens(predicted_tokens)
-
-    # Filter out special tokens and join the remaining tokens into a string
-    extracted_text = " ".join([token for token in tokens if token not in layoutlm_processor.tokenizer.all_special_tokens])
+        outputs = qa_model.generate(inputs, max_length=150, num_return_sequences=1)
 
-    return extracted_text
+    # Decode the response
+    response = qa_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
 # Streamlit App
 st.set_page_config(page_title="Invoice Extractor")
@@ -43,6 +49,7 @@ if submit:
     if image is None:
         st.warning("Please upload an image.")
     else:
-        extracted_text = extract_text_from_image(image, question)
+        extracted_text = extract_text_from_image(image)
+        response = get_response_from_llm(extracted_text, question)
         st.subheader("Extracted Information:")
-        st.write(extracted_text)
+        st.write(response)