Update app.py

app.py

@@ -7,12 +7,22 @@ from PIL import Image
 from transformers import AutoProcessor, AutoModelForTokenClassification
 
 #function to load gemini pro vision model and get response
-def get_model_response(input,image,prompt):  #input: how LLM model behave like, #image: To extract info, prompt: ask something
-    #load a model
-    processor = AutoProcessor.from_pretrained("monuirctc/invoice-extraction")
-    model = AutoModelForTokenClassification.from_pretrained("monuirctc/invoice-extraction")
-    response=model.generate_content([input,image[0],prompt]) #get response from model
-    return response.text
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+
+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)
+
+    # Tokenize inputs
+    inputs = tokenizer(input_text, image, prompt, return_tensors="pt", max_length=512, truncation=True)
+
+    # Generate response from the model
+    outputs = model(**inputs)
+    response = tokenizer.decode(outputs.logits.argmax(dim=-1))
+
+    return response
 
 #function to provide image
 def input_image_setup(uploaded_file):