generation

app.py

@@ -4,25 +4,28 @@ import os
 import gradio as gr
 from huggingface_hub import snapshot_download
 import spaces
+import torch
+from transformers import T5TokenizerFast
+from accelerate import Accelerator
 
 from pix2struct.modeling import Pix2StructModel
 from pix2struct.processing import extract_patches
-from pix2struct.inference import ask_generator, generate
+from pix2struct.inference import ask_generator, generate, DocumentQueries, DocumentQuery
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 hub_token = os.environ.get('HUB_TOKEN')
 model_path = snapshot_download('artyomxyz/pix2struct-docmatix', use_auth_token=hub_token)
-model = Pix2StructModel.load(model_path)
-model.eval()
 
 tokenizer = T5TokenizerFast.from_pretrained('google/pix2struct-base')
 
 @spaces.GPU
-def ask(image, question):
+def ask(image, questions):
+    questions = questions.split('\n')
     accelerator = Accelerator(mixed_precision="bf16")
+    model = Pix2StructModel.load(model_path)
+    model.eval()
     model = accelerator.prepare(model)
-
     documents = [
         DocumentQueries(
             meta=None,
@@ -30,21 +33,22 @@ def ask(image, question):
             queries=[
                 DocumentQuery(
                     meta=None,
-                    generator=ask_generator(tokenizer, qa['question'])
+                    generator=ask_generator(tokenizer, question)
                 )
+                for question in questions
             ]
         )
     ]
-    result = generate(model, documents, device=accelerator.device)
-
-
-    return result[0].queries[0].output
+    with torch.inference_mode():
+        with accelerator.autocast():
+            result = generate(model, documents, device=accelerator.device)
+    return '\n'.join([q.output for q in result[0].queries])
 
 demo = gr.Interface(
     fn=ask,
     inputs=[
         gr.Image(type='numpy'),
-        gr.Textbox(),
+        gr.Textbox(label="Questions (one question per line)"),
     ],
     outputs='text'
 )