Update app.py

app.py

@@ -1,27 +1,10 @@
 from PIL import Image
 import streamlit as st
-from transformers import AutoModelForImageToText, AutoTokenizer, AutoModelForCausalLM, AutoTokenizer
-import torch
+from transformers import LayoutLMForTokenClassification, LayoutLMTokenizer
 
 # Load Hugging Face models
-image_to_text_model = AutoModelForImageToText.from_pretrained("microsoft/layoutlm-base-uncased")
-image_to_text_tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
-text_generation_model = AutoModelForCausalLM.from_pretrained("gpt2")
-text_generation_tokenizer = AutoTokenizer.from_pretrained("gpt2")
-
-# Function to extract text from image
-def extract_text_from_image(image):
-    inputs = image_to_text_tokenizer(image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = image_to_text_model(**inputs)
-    return image_to_text_tokenizer.batch_decode(outputs.logits, skip_special_tokens=True)[0]
-
-# Function to generate response
-def generate_response(prompt, text):
-    inputs = text_generation_tokenizer(prompt + text, return_tensors="pt")
-    with torch.no_grad():
-        outputs = text_generation_model.generate(**inputs)
-    return text_generation_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+layoutlm_model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased")
+layoutlm_tokenizer = LayoutLMTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
 
 # Streamlit App
 st.set_page_config(page_title="Invoice Extractor")
@@ -42,3 +25,19 @@ if uploaded_file is not None:
             st.write(response)
         else:
             st.warning("Please upload an image.")
+
+# Function to extract text from image
+def extract_text_from_image(image):
+    inputs = layoutlm_tokenizer(image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = layoutlm_model(**inputs)
+    # Assume we want to extract only the text from the first token
+    tokens = layoutlm_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+    text = " ".join([token for token in tokens if token != "[PAD]"])
+    return text
+
+# Function to generate response
+def generate_response(prompt, text):
+    # You can use your existing code here for response generation
+    # For this example, let's simply return the extracted text with the prompt
+    return f"{prompt} {text}"