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)