Update app.py

app.py

@@ -2,14 +2,86 @@ import gradio as gr
 # from PIL import Image
 from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 
-model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-docvqa-large")
-processor = Pix2StructProcessor.from_pretrained("google/pix2struct-docvqa-large")
+def _add_markup(table):
+    parts = [p.strip() for p in table.splitlines(keepends=False)]
+    if parts[0].startswith('TITLE'):
+        result = f"Title: {parts[0].split(' | ')[1].strip()}\n"
+        rows = parts[1:]
+    else:
+        result = ''
+        rows = parts
+    prefixes = ['Header: '] + [f'Row {i+1}: ' for i in range(len(rows) - 1)]
+    return result + '\n'.join(prefix + row for prefix, row in zip(prefixes, rows))
+
+
+_TABLE = """Year | Democrats | Republicans | Independents
+2004 | 68.1% | 45.0% | 53.0%
+2006 | 58.0% | 42.0% | 53.0%
+2007 | 59.0% | 38.0% | 45.0%
+2009 | 72.0% | 49.0% | 60.0%
+2011 | 71.0% | 51.2% | 58.0%
+2012 | 70.0% | 48.0% | 53.0%
+2013 | 72.0% | 41.0% | 60.0%"""
+
+_INSTRUCTION = 'Read the table below to answer the following questions.'
+
+
+_TEMPLATE = f"""{_INSTRUCTION}
+{_add_markup(_TABLE)}
+Q: In which year republicans have the lowest favor rate?
+A: Let's find the column of republicans. Then let's extract the favor rates, they [45.0, 42.0, 38.0, 49.0, 51.2, 48.0, 41.0]. The smallest number is 38.0, that's Row 3.  Row 3 is year 2007. The answer is 2007.
+Q: What is the sum of Democrats' favor rates of 2004, 2012, and 2013?
+A: Let's find the rows of years 2004, 2012, and 2013. We find Row 1, 6, 7. The favor dates of Demoncrats on that 3 rows are 68.1, 70.0, and 72.0. 68.1+70.0+72=210.1. The answer is 210.1.
+Q: By how many points do Independents surpass Republicans in the year of 2011?
+A: Let's find the row with year = 2011. We find Row 5. We extract Independents and Republicans' numbers. They are 58.0 and 51.2. 58.0-51.2=6.8. The answer is 6.8.
+Q: Which group has the overall worst performance?
+A: Let's sample a couple of years. In Row 1, year 2004, we find Republicans having the lowest favor rate 45.0 (since 45.0<68.1, 45.0<53.0). In year 2006, Row 2, we find Republicans having the lowest favor rate 42.0 (42.0<58.0, 42.0<53.0). The trend continues to other years. The answer is Republicans.
+Q: Which party has the second highest favor rates in 2007?
+A: Let's find the row of year 2007, that's Row 3. Let's extract the numbers on Row 3: [59.0, 38.0, 45.0]. 45.0 is the second highest. 45.0 is the number of Independents. The answer is Independents.
+{_INSTRUCTION}"""
+
+def text_generate(prompt, table, problem):
+    p = prompt + "\n" + _INSTRUCTION + "\n" + table + "\n" + "Q: " + problem
+    # print(f"Final prompt is : {p}")
+    json_ = {"inputs": p,
+            "parameters":
+            {
+            "top_p": 0.9,
+          "temperature": 1.1,
+          "max_new_tokens": 64,
+          "return_full_text": True
+          }, "options": 
+              {
+              "use_cache": True,
+              "wait_for_model":True
+              },}
+    response = requests.post(API_URL, headers=headers, json=json_)
+    print(f"Response  is : {response}")
+    output = response.json()
+    print(f"output is : {output}") #{output}")
+    output_tmp = output[0]['generated_text']
+    print(f"output_tmp is: {output_tmp}")
+    #solution = output_tmp.split("\nQ:")[0]   #output[0]['generated_text'].split("Q:")[0] # +"."
+    #print(f"Final response after splits is: {solution}")
+
+    #return solution 
+    return output_tmp
+
+
+
+
+model_deplot = Pix2StructForConditionalGeneration.from_pretrained("belkada/deplot")
+processor_deplot = Pix2StructProcessor.from_pretrained("belkada/deplot")
 
 def process_document(image, question):
     # image = Image.open(image)
-    inputs = processor(images=image, text=question, return_tensors="pt")
-    predictions = model.generate(**inputs)
-    return processor.decode(predictions[0], skip_special_tokens=True)
+    inputs = processor_deplot(images=image, text="Generate the underlying data table for the figure below:", return_tensors="pt")
+    predictions = model_deplot.generate(**inputs)
+    table = processor_deplot.decode(predictions[0], skip_special_tokens=True)
+
+    # send prompt+table to LLM
+    res = text_generate(_TEMPLATE, table, question)
+    print (res)
  
 description = "Demo for pix2struct fine-tuned on DocVQA (document visual question answering). To use it, simply upload your image and type a question and click 'submit', or click one of the examples to load them. Read more at the links below."
 article = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2210.03347.pdf' target='_blank'>PIX2STRUCT: SCREENSHOT PARSING AS PRETRAINING FOR VISUAL LANGUAGE UNDERSTANDING</a></p>"
@@ -18,7 +90,7 @@ demo = gr.Interface(
     fn=process_document,
     inputs=["image", "text"],
     outputs="text",
-    title="Demo: pix2struct for DocVQA",
+    title="Demo: deplot+llm test",
     description=description,
     article=article,
     enable_queue=True,