Aeon-Avinash commited on
Commit
b2c2b10
1 Parent(s): 8803f27

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +162 -0
app.py ADDED
@@ -0,0 +1,162 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ import torch
3
+ import gradio as gr
4
+ import os
5
+ from PIL import Image
6
+ import pytesseract
7
+ import tempfile
8
+ import shutil
9
+ from pdf2image import convert_from_path
10
+
11
+ model_name = "deepset/roberta-base-squad2"
12
+ text_qna = pipeline("question-answering", model=model_name, tokenizer=model_name)
13
+ vision_qna = pipeline("document-question-answering", model="impira/layoutlm-document-qa")
14
+
15
+ # Vision QnA requires: PyTesseract for OCR. Tesseract executable needs to be installed separately.
16
+ # sudo apt install tesseract-ocr (https://tesseract-ocr.github.io/tessdoc/Installation.html)
17
+
18
+ def load_file(file_input, encoding = 'utf-8'):
19
+ if not os.path.exists(file_input):
20
+ raise FileNotFoundError(f"The file does not exist.")
21
+
22
+ with open(file_input, 'r', encoding=encoding) as file:
23
+ try:
24
+ content = file.read()
25
+ except UnicodeDecodeError:
26
+ # If a UnicodeDecodeError occurs, try reading with 'latin1' encoding
27
+ with open(file_input, 'r', encoding='latin1') as file:
28
+ content = file.read()
29
+
30
+ return content
31
+
32
+ def save_image(file):
33
+ try:
34
+ temp_dir = tempfile.mkdtemp()
35
+ file_path = os.path.join(temp_dir, os.path.basename(file.name))
36
+ # Copy the file from the temporary Gradio directory to our temporary directory
37
+ shutil.copyfile(file.name, file_path)
38
+ # when working with saving image files through Gradio,
39
+ # using `shutil.copyfile` to handle `NamedString` objects for file uploads is the correct approach
40
+
41
+ return file_path
42
+ except Exception as e:
43
+ print(e)
44
+
45
+
46
+ def save_pdf(file):
47
+ temp_dir = tempfile.mkdtemp()
48
+ pdf_path = os.path.join(temp_dir, os.path.basename(file.name))
49
+
50
+ # Copy the file from the temporary Gradio directory to our temporary directory
51
+ shutil.copyfile(file.name, pdf_path)
52
+
53
+ # Convert PDF to images
54
+ images = convert_from_path(pdf_path)
55
+
56
+ image_paths = []
57
+ for i, img in enumerate(images):
58
+ image_path = os.path.join(temp_dir, f'page_{i}.png')
59
+ img.save(image_path, 'PNG')
60
+ image_paths.append(image_path)
61
+
62
+ print(image_paths)
63
+ return image_paths
64
+
65
+
66
+ def qna_text_content(content, question):
67
+ result = text_qna(question=question, context=content)
68
+ return result
69
+
70
+ def qna_image_content(content, question):
71
+ # result = vision_qna(question=question, image=content)
72
+ result = vision_qna(content, question)
73
+ print(f"image question: {question}")
74
+ return result
75
+
76
+ def qna_pdf_content(image_paths, question):
77
+ answers = []
78
+ try:
79
+ for image_path in image_paths:
80
+ result = vision_qna(image=image_path, question=question)
81
+ print(result[0]['answer'], result[0]['score'])
82
+ answers.append(result[0]['answer'])
83
+ return " \n".join(answers)
84
+ except Exception as e:
85
+ return f"An error occurred during processing: {e}"
86
+
87
+ def answer_the_question_for_doc(text_input, file_input, question):
88
+ # Order of input parameters is Imp. for Gradio to accept respective Inputs
89
+ if file_input is not None:
90
+ print(f"File type: {type(file_input)}")
91
+ print(f"File name: {file_input.name}")
92
+
93
+ file_extension = file_input.name.split('.')[-1].lower()
94
+ if file_extension in ['txt']:
95
+ try:
96
+ content = load_file(file_input)
97
+ if not content or not question:
98
+ return "Please provide both content and a question."
99
+ result = qna_text_content(content, question)
100
+ return result["answer"]
101
+
102
+ except FileNotFoundError or Exception as e:
103
+ print(e)
104
+ exit(1)
105
+
106
+ elif file_extension in ['png', 'jpeg', 'jpg']:
107
+ try:
108
+ img_file_path = save_image(file_input)
109
+
110
+ if not question:
111
+ return "Please provide a question."
112
+ result = qna_image_content(img_file_path, question)
113
+ print(result)
114
+ return result[0]["answer"]
115
+
116
+ except Exception as e:
117
+ return f"An error occurred during vision processing: {e}"
118
+
119
+ elif file_extension in ['pdf']:
120
+ try:
121
+ image_paths = save_pdf(file_input)
122
+
123
+ if not question:
124
+ return "Please provide a question."
125
+ result = qna_pdf_content(image_paths, question)
126
+ print(result)
127
+ return result
128
+
129
+ except Exception as e:
130
+ return f"An error occurred during vision processing: {e}"
131
+
132
+ else:
133
+ return "Unsupported file type. Please upload a .txt, ,.pdf, .png, or .jpeg file."
134
+ else:
135
+ if not text_input or not question:
136
+ return "Please provide both content and a question."
137
+ content = text_input
138
+ result = qna_text_content(content, question)
139
+ return result["answer"]
140
+
141
+ gr.close_all()
142
+
143
+ with gr.Blocks() as demo:
144
+ gr.Markdown("# QnA System")
145
+ gr.Markdown("This App answers a question based on text content or uploaded file (txt, png, jpeg, pdf).")
146
+
147
+ with gr.Row():
148
+ text_input = gr.Textbox(label="Text Input", placeholder="Enter text content here...")
149
+ file_input = gr.File(label="File Upload", file_types=['txt', 'png', 'jpeg', 'pdf'])
150
+
151
+ question = gr.Textbox(label="Question", placeholder="Enter your question here...")
152
+ output = gr.Textbox(label="Answer", placeholder="The answer will appear here...")
153
+
154
+ text_input.change(lambda x: gr.update(visible=not x), inputs=text_input, outputs=file_input)
155
+ file_input.change(lambda x: gr.update(visible=not x), inputs=file_input, outputs=text_input)
156
+
157
+ button = gr.Button("Get Answer")
158
+ button.click(answer_the_question_for_doc, inputs=[text_input, file_input, question],
159
+ outputs=output)
160
+ demo.launch()
161
+
162
+ # print(qna_image_content("https://gradientflow.com/wp-content/uploads/2023/10/newsletter87-RAG-simple.png", "What is the step prior to embedding?"))