Skip to content

Saroj Dhiman / invoice information extractor

Go to a project
Commit 69c92f3b by Saroj Dhiman
Options

Update model_2.py

1 parent 3e4cce28
Inline Side-by-side
Showing with 2 additions and 2 deletions
model_2.py
......@@ -11,8 +11,8 @@ import time
DEVICE = "cpu"
# DEVICE = "cuda:0"
# Load the Pix2Struct model and processor
model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-docvqa-base").to(DEVICE)
processor = Pix2StructProcessor.from_pretrained("google/pix2struct-docvqa-base")
model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-docvqa-large").to(DEVICE)
processor = Pix2StructProcessor.from_pretrained("google/pix2struct-docvqa-large")
def generate(img, questions):
inputs = processor(images=[img for _ in range(len(questions))],
......
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!