Update app.py

app.py

@@ -3,9 +3,11 @@ import re
 from PIL import Image
 import requests
 from nougat.dataset.rasterize import rasterize_paper
-
 from transformers import NougatProcessor, VisionEncoderDecoderModel
 import torch
+import gradio as gr
+import uuid
+import os
 
 processor = NougatProcessor.from_pretrained("facebook/nougat-small")
 model = VisionEncoderDecoderModel.from_pretrained("facebook/nougat-small")
@@ -48,7 +50,7 @@ def predict(image):
 
 
 
-def inference(pdf_file, pdf_link):
+def inference(pdf_file, pdf_link, file_btn):
   if pdf_file is None:
     if pdf_link == '':
       print("No file is uploaded and No link is provided")
@@ -61,19 +63,20 @@ def inference(pdf_file, pdf_link):
 
   images = rasterize_paper(file_name, return_pil=True)
   sequence = ""
-  # infer for every page and concat
+  # infer for every page and concat
   for image in images:
     sequence += predict(image)
 
 
   content = sequence.replace(r'\(', '$').replace(r'\)', '$').replace(r'\[', '$$').replace(r'\]', '$$')
-  return content
 
-import gradio as gr
-import uuid
-import os
-import requests
-import re
+  if file_btn:
+    with open(f"{os.getcwd()}/output.txt","w+") as f:
+      f.write(content)
+      f.close()
+    
+  return content, f"{os.getcwd()}/output.txt"
+
 
 css = """
   #mkd {
@@ -98,25 +101,30 @@ with gr.Blocks(css=css) as demo:
   with gr.Row(equal_height=True):
     pdf_file = gr.File(label='PDF 📑', file_count='single', scale=1)
     pdf_link = gr.Textbox(placeholder='Enter an arxiv link here', label='Link to Paper🔗', scale=1)
-
+  with gr.Row():
+    file_btn = gr.Checkbox(label='Download output as file 📑')
   with gr.Row():
     btn = gr.Button('Run Nougat 🍫')
-    clr = gr.Button('Clear 🧼')
+  with gr.Row():
+    clr = gr.Button('Clear Inputs & Outputs 🧼')
 
-  output_headline = gr.Markdown("PDF converted to markup language through Nougat-OCR👇")
-  parsed_output = gr.Markdown(elem_id='mkd', value='OCR Output 📝')
+  output_headline = gr.Markdown("## PDF converted to markup language through Nougat-OCR👇")
+  with gr.Row():
+      parsed_output = gr.Markdown(elem_id='mkd', value='Output Text 📝')
+      output_file = gr.File(file_types = ["txt"], label="Output File 📑")
   
-  btn.click(inference, [pdf_file, pdf_link], parsed_output )
+  btn.click(inference, [pdf_file, pdf_link, file_btn], [parsed_output, output_file])
   clr.click(lambda : (gr.update(value=None), 
                       gr.update(value=None),
+                      gr.update(value=None), 
                       gr.update(value=None)), 
              [], 
-             [pdf_file, pdf_link, parsed_output]
+             [pdf_file, pdf_link, file_btn, parsed_output, output_file]
             )
   gr.Examples(
-      [["nougat.pdf", ""], [None, "https://arxiv.org/pdf/2308.08316.pdf"]],
-      inputs = [pdf_file, pdf_link],
-      outputs = parsed_output,
+      [["nougat.pdf", "", True], [None, "https://arxiv.org/pdf/2308.08316.pdf", True]],
+      inputs = [pdf_file, pdf_link, file_btn],
+      outputs = [parsed_output, output_file],
       fn=inference,
       cache_examples=True,
       label='Click on any Examples below to get Nougat OCR results quickly:'