File size: 5,709 Bytes
77eac00 d15e89e 77eac00 e5c702c 8ae9c70 77eac00 e602714 77eac00 e5c702c 23b7c2e e5c702c 77eac00 d15e89e 1482457 d15e89e 77eac00 2d5d236 23c0f1b 77eac00 c3edca2 9fc3c5a caccbf6 9fc3c5a d15e89e 9fc3c5a caccbf6 9fc3c5a 77eac00 9fc3c5a 77eac00 9fc3c5a e5c702c c3edca2 9fc3c5a caccbf6 4fc375b 6d6a410 4fc375b caccbf6 c3edca2 4dff3e2 c3edca2 3fa8524 c3edca2 4dff3e2 c3edca2 caccbf6 c3edca2 4dff3e2 c3edca2 77eac00 caccbf6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import gradio as gr
from time import time
from pathlib import Path
from gradio_pdf import PDF
from pdf2image import convert_from_path
import shutil
import tempfile
from transformers import pipeline
import subprocess as sp
out_files = gr.State([])
FILE_TIMEOUT = 10 ** 3
MAX_FILES = 10
p = pipeline(
"document-question-answering",
model="TusharGoel/LiLT-Document-QA",
)
custom_css = """
<style>
button.primary {
background-color: #2AD6B5 !important;
}
</style>
"""
def handle_files(cur_files):
cur_time = cur_files[-1][0]
deleted_indices = set()
for other_idx, (other_time, other_file) in enumerate(cur_files[:-1]):
if abs(cur_time - other_time) > FILE_TIMEOUT:
shutil.rmtree(other_file.parent)
deleted_indices.add(other_idx)
cur_files = [cur_files[idx] for idx in range(len(cur_files)) if idx not in deleted_indices]
if len(cur_files) > MAX_FILES:
for _, other_file in cur_files[:-MAX_FILES]:
shutil.rmtree(other_file.parent)
cur_files = cur_files[-MAX_FILES:]
return cur_files
def extract_text(pdf_file):
"""
Generate a text rendering of a PDF file in the form of a list of lines.
"""
args = ['pdftotext', '-layout', pdf_file, '-']
cp = sp.run(
args, stdout=sp.PIPE, stderr=sp.DEVNULL,
check=True, text=True
)
return cp.stdout
# Function to process PDF and generate ZIP file
def process_pdf(pdf_file, cur_files):
zip_output = Path(tempfile.mkdtemp()) / f'{Path(pdf_file).stem}'
# zip_output.parent.mkdir()
with tempfile.TemporaryDirectory() as path:
pdf_output = path
convert_from_path(pdf_file, output_folder=str(pdf_output))
# Create a BytesIO object to store zip file in memory
shutil.make_archive(zip_output, 'zip', pdf_output)
zip_output = zip_output.with_suffix('.zip')
cur_time = time()
cur_files.append((cur_time, zip_output))
cur_files = handle_files(cur_files)
return str(zip_output), cur_files
def interact_with_pdf(doc, question):
with tempfile.TemporaryDirectory() as path:
images = convert_from_path(doc, output_folder=path)
outputs = []
for img in images:
outputs += p(img, question)
return sorted(outputs, key=lambda x: x["score"], reverse=True)[0]['answer']
'''
text_interface = gr.Interface(
fn=extract_text,
inputs=PDF(label="Загрузить PDF"),
outputs=gr.Textbox(label="Полученный текст"),
title="PDF в текст",
description="Сервис вынимает содержащийся в PDF контейнере текст."
)
pdf_interface = gr.Interface(
fn=process_pdf,
inputs=[PDF(label="Загрузить PDF"), out_files],
outputs=[gr.File(label="Скачать ZIP архив"), out_files],
title="PDF в картинки",
description="Переводит постранично файл в изображения, и предлагает бскачать в виде ZIP архива."
)
image_interface = gr.Interface(
fn=interact_with_pdf,
inputs=[
PDF(label="Загрузить PDF"),
gr.Textbox(label="Что хотите найти?")
],
outputs=gr.Textbox(label="Возможный ответ"),
title="Спроси PDF",
description="Сервис через методы OCR сканирует загруженный файл для нахождения ответа на заданный вопрос."
)
# Create a tabbed interface
tabbed_interface = gr.TabbedInterface(
[text_interface, pdf_interface, image_interface],
title="Взаимодействие с PDF",
tab_names=["В текст", "В картинки", "Задай вопрос"],
# description="Choose a tab to perform the desired task.",
css=custom_css,
analytics_enabled=True
)
'''
theme = gr.themes.Default(primary_hue="blue").set(
loader_color="#2AD6B5",
slider_color="#2AD6B5",
button_primary_background_fill="#2AD6B5",
button_primary_text_color="#000000",
button_primary_background_fill_hover="#20D0B0",
button_primary_text_color_hover="#000000",
button_primary_border_color="999999",
button_primary_border_color_hover="333333"
)
with gr.Blocks(theme=theme) as tabbed_interface:
with gr.Tab("В текст"):
with gr.Row():
with gr.Column():
inp1 = PDF(label="Загрузить PDF")
btn_sbm1 = gr.Button("Запустить", variant="primary")
with gr.Column():
out1 = gr.Textbox(label="Полученный текст")
btn_sbm1.click(fn=extract_text, inputs=inp1, outputs=out1)
with gr.Tab("В картинки"):
with gr.Row():
with gr.Column():
inp2 = PDF(label="Загрузить PDF")
btn_sbm2 = gr.Button("Запустить", variant="primary")
with gr.Column():
out2 = gr.File(label="Скачать ZIP архив")
btn_sbm2.click(fn=process_pdf, inputs=[inp2, out_files], outputs=[out2, out_files])
with gr.Tab("Задай вопрос"):
with gr.Row():
with gr.Column():
inp3 = PDF(label="Загрузить PDF")
inp4 = gr.Textbox(label="Что хотите найти?")
btn_sbm3 = gr.Button("Запустить", variant="primary")
with gr.Column():
out3 = gr.Textbox(label="Возможный ответ")
btn_sbm3.click(fn=interact_with_pdf, inputs=[inp3, inp4], outputs=out3)
tabbed_interface.launch(show_api=False, max_threads=8)
|