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")
@@ -67,13 +69,14 @@ def inference(pdf_file, pdf_link):
 
 
   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("output.txt","w+") as f:
+      f.write(content)
+      f.close()
+    
+  return content, "output.txt"
+
 
 css = """
   #mkd {
@@ -98,20 +101,24 @@ 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👇")
+  output_headline = gr.Markdown("## PDF converted to markup language through Nougat-OCR👇")
   parsed_output = gr.Markdown(elem_id='mkd', value='OCR Output 📝')
+  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, output_file, parsed_output]
             )
   gr.Examples(
       [["nougat.pdf", ""], [None, "https://arxiv.org/pdf/2308.08316.pdf"]],