Update app.py

app.py

@@ -4,16 +4,16 @@
 import os
 import streamlit as st
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForTokenClassification
 
-#function to load gemini pro vision model and get response
-from transformers import AutoTokenizer, AutoModelForTokenClassification
+# Load model directly
+from transformers import AutoProcessor, AutoModelForSeq2SeqLM
+
+
 
 def get_model_response(input_text, image, prompt):
-    # Load the pre-trained model and tokenizer
-    model_name = "monuirctc/invoice-extraction"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForTokenClassification.from_pretrained(model_name)
+    # Load the pre-trained model
+    processor = AutoProcessor.from_pretrained("google/pix2struct-base")
+    model = AutoModelForSeq2SeqLM.from_pretrained("google/pix2struct-base")
 
     # Tokenize inputs
     inputs = tokenizer(input_text, image, prompt, return_tensors="pt", max_length=512, truncation=True)