| from transformers import pipeline |
| import torch |
| import gradio as gr |
| import os |
| from PIL import Image |
| import pytesseract |
| import tempfile |
| import shutil |
| from pdf2image import convert_from_path |
| import json |
| import spaces |
|
|
| model_name = "deepset/roberta-base-squad2" |
|
|
| |
| |
|
|
| vision_qna = pipeline("document-question-answering", model="impira/layoutlm-document-qa", device=0) |
|
|
| def load_file(file_input, encoding = 'utf-8'): |
| if not os.path.exists(file_input): |
| raise FileNotFoundError(f"The file does not exist.") |
|
|
| with open(file_input, 'r', encoding=encoding) as file: |
| try: |
| content = file.read() |
| except UnicodeDecodeError: |
| |
| with open(file_input, 'r', encoding='latin1') as file: |
| content = file.read() |
|
|
| return content |
|
|
| def save_image(file): |
| try: |
| temp_dir = tempfile.mkdtemp() |
| file_path = os.path.join(temp_dir, os.path.basename(file.name)) |
| |
| shutil.copyfile(file.name, file_path) |
| |
| |
|
|
| return file_path |
| except Exception as e: |
| print(e) |
|
|
|
|
| def save_pdf(file): |
| temp_dir = tempfile.mkdtemp() |
| pdf_path = os.path.join(temp_dir, os.path.basename(file.name)) |
| |
| |
| shutil.copyfile(file.name, pdf_path) |
| |
| |
| images = convert_from_path(pdf_path) |
| |
| image_paths = [] |
| for i, img in enumerate(images): |
| image_path = os.path.join(temp_dir, f'page_{i}.png') |
| img.save(image_path, 'PNG') |
| image_paths.append(image_path) |
| |
| |
| return image_paths |
|
|
|
|
| def qna_image_content(content, question): |
| result = vision_qna(content, question) |
| page_result = get_page_result(result, question) |
| return get_result_json([page_result]) |
| |
| |
| |
|
|
| def get_page_result(raw_result, question): |
| |
| page_result = { |
| "result": { |
| "question": question, |
| "answer": raw_result[0]['answer'] |
| }, |
| "confidence": raw_result[0]['score'] |
| } |
| return page_result |
| |
| def get_result_json(page_results): |
| result = { |
| "pages": page_results |
| } |
| return json.dumps(result) |
| def qna_pdf_content(image_paths, question, short_circuit_threshold): |
| all_page_results = [] |
| try: |
| for image_path in image_paths: |
| result = vision_qna(image=image_path, question=question) |
| |
| page_result = get_page_result(result, question) |
| all_page_results.append(page_result) |
| if result[0]['score'] >= short_circuit_threshold: |
| break |
| return get_result_json(all_page_results) |
| except Exception as e: |
| return f"An error occurred during processing: {e}" |
|
|
| @spaces.GPU |
| def predict(doc, question, short_circuit_threshold): |
| |
| file_input = doc |
| if file_input is not None and question is not None: |
| print(f"File type: {type(file_input)}") |
| print(f"File name: {file_input.name}") |
|
|
| filename_lower = file_input.name.lower() |
| is_pdf = filename_lower.endswith('.pdf') or filename_lower.endswith("blob") |
| if is_pdf: |
| try: |
| image_paths = save_pdf(file_input) |
| |
| if not question: |
| return "Please provide a question." |
| result = qna_pdf_content(image_paths, question, short_circuit_threshold) |
| |
| return result |
| except Exception as e: |
| return f"An error occurred during vision processing: {e}" |
|
|
| else: |
| try: |
| img_file_path = save_image(file_input) |
| |
| if not question: |
| return "Please provide a question." |
| result = qna_image_content(img_file_path, question) |
| |
| return result |
|
|
| except Exception as e: |
| return f"An error occurred during vision processing: {e}" |
|
|
| else: |
| raise gr.Error("Provide both a document and a question") |
|
|
| gr.close_all() |
|
|
| with gr.Blocks() as demo: |
| gr.Markdown("# LayoutLM Document QnA") |
| gr.Markdown("This app performs visual document QnA for docs in PDF or image formats.") |
|
|
| file_input = gr.File(label="File Upload", file_types=['png', 'jpeg', 'jpg', 'bmp', 'gif', 'pdf']) |
|
|
| question = gr.Textbox(label="Question", placeholder="Enter your question here...") |
| threshold_input = gr.Number(label="Short-circuit threshold", value=0.97, minimum=0, maximum=1) |
| output = gr.JSON(label="Answer") |
|
|
|
|
| button = gr.Button("Run model") |
| button.click(predict, inputs=[file_input, question, threshold_input], |
| outputs=output) |
| demo.launch() |
|
|
| |
