Update app.py

app.py

@@ -1,151 +1,42 @@
 import gradio as gr
-import csv
-from io import StringIO
-from PIL import Image
-import numpy as np
-import base64
-
-# Define the annotation types
-ANNOTATION_TYPES = ['rect', 'circle']
-
-# Define the Annotation class
-class Annotation:
-    def __init__(self, x, y, width, height, annotation_type):
-        self.x = x
-        self.y = y
-        self.width = width
-        self.height = height
-        self.type = annotation_type
-
-# Define the Gradio interface
-def annotate_images(images):
-    # Define the canvas size
-    canvas_size = (600, 600)
-
-    # Define the initial state
-    state = {
-        'image': None,
-        'annotations': [],
-        'annotation_type': ANNOTATION_TYPES[0],
-        'start_point': None
-    }
-
-    # Define the canvas drawing function
-    def draw_canvas(canvas, image_data, annotations):
-        # Convert the image data to a PIL Image object
-        image = Image.fromarray(image_data)
-
-        # Resize the image to fit the canvas
-        image = image.resize(canvas_size)
-
-        # Draw the image on the canvas
-        canvas.draw_image(image, (canvas_size[0]/2, canvas_size[1]/2))
-
-        # Draw the annotations on the canvas
-        for annotation in annotations:
-            x, y, width, height = annotation.x, annotation.y, annotation.width, annotation.height
-            if annotation.type == 'rect':
-                canvas.draw_rect(x, y, width, height, stroke_color='red')
-            elif annotation.type == 'circle':
-                radius = np.sqrt(np.power(width, 2) + np.power(height, 2)) / 2
-                center_x, center_y = x + width / 2, y + height / 2
-                canvas.draw_circle(center_x, center_y, radius, stroke_color='red')
-
-    # Define the canvas mousedown event handler
-    def canvas_mousedown(canvas, x, y):
-        state['start_point'] = (x, y)
-
-    # Define the canvas mousemove event handler
-    def canvas_mousemove(canvas, x, y):
-        if state['start_point'] is not None:
-            start_x, start_y = state['start_point']
-            end_x, end_y = x, y
-            annotation_type = state['annotation_type']
-            draw_annotation(canvas, start_x, start_y, end_x, end_y, annotation_type)
-
-    # Define the canvas mouseup event handler
-    def canvas_mouseup(canvas, x, y):
-        if state['start_point'] is not None:
-            start_x, start_y = state['start_point']
-            end_x, end_y = x, y
-            annotation_type = state['annotation_type']
-            add_annotation(start_x, start_y, end_x, end_y, annotation_type)
-            state['start_point'] = None
-
-    # Define the add annotation function
-    def add_annotation(start_x, start_y, end_x, end_y, annotation_type):
-        # Calculate the width and height of the annotation
-        width = np.abs(start_x - end_x)
-        height = np.abs(start_y - end_y)
-
-        # Create the annotation object
-        annotation = Annotation(start_x, start_y, width, height, annotation_type)
-
-        # Add the annotation to the array
-        state['annotations'].append(annotation)
-
-        # Redraw the canvas
-        draw_canvas(canvas, state['image'], state['annotations'])
-
-    #
-
-            # Define the draw annotation function
-    def draw_annotation(canvas, start_x, start_y, end_x, end_y, annotation_type):
-        canvas.clear()
-        draw_canvas(canvas, state['image'], state['annotations'])
-        width = np.abs(start_x - end_x)
-        height = np.abs(start_y - end_y)
-        if annotation_type == 'rect':
-            canvas.draw_rect(start_x, start_y, width, height, stroke_color='red')
-        elif annotation_type == 'circle':
-            radius = np.sqrt(np.power(width, 2) + np.power(height, 2)) / 2
-            center_x, center_y = start_x + width / 2, start_y + height / 2
-            canvas.draw_circle(center_x, center_y, radius, stroke_color='red')
-
-    # Define the annotation type dropdown event handler
-    def annotation_type_changed(value):
-        state['annotation_type'] = value
-
-    # Define the download annotations button click event handler
-    def download_annotations_clicked():
-        # Define the csv headers
-        headers = ['x', 'y', 'width', 'height', 'type']
-
-        # Define the csv data
-        rows = [[str(annotation.x), str(annotation.y), str(annotation.width), str(annotation.height), annotation.type]
-                for annotation in state['annotations']]
-
-        # Create the csv string
-        csv_string = StringIO()
-        csv_writer = csv.writer(csv_string)
-        csv_writer.writerow(headers)
-        for row in rows:
-            csv_writer.writerow(row)
-
-        # Download the csv file
-        b64_csv = base64.b64encode(csv_string.getvalue().encode()).decode()
-        href = f'data:text/csv;base64,{b64_csv}'
-        download_link = f'<a href="{href}" download="annotations.csv">Download Annotations CSV</a>'
-        gr.Interface.show(download_link)
-
-    # Define the interface components
-    image = gr.inputs.Image(label='Image')
-    annotation_type = gr.inputs.Dropdown(ANNOTATION_TYPES, label='Annotation Type', default=ANNOTATION_TYPES[0], onchange=annotation_type_changed)
-    download_annotations = gr.outputs.Button(label='Download Annotations', type='button', onclick=download_annotations_clicked)
-    canvas = gr.outputs.Canvas(draw_event_handlers={
-        'mousedown': canvas_mousedown,
-        'mousemove': canvas_mousemove,
-        'mouseup': canvas_mouseup
-    })
-
-    # Define the interface function
-    def annotate_images(images):
-        state['image'] = images[0]
-        draw_canvas(canvas, state['image'], state['annotations'])
-        return canvas, annotation_type, download_annotations
-
-    # Create the interface
-    interface = gr.Interface(annotate_images, inputs=image, outputs=[canvas, annotation_type, download_annotations], capture_session=True)
-
-    return interface
-
+import requests
+import io
+import json
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+
+# Download and load pre-trained model and tokenizer
+model_name = "distilbert-base-cased-distilled-squad"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+
+def answer_question(pdf_file, question):
+    # Convert PDF to text
+    pdf_data = pdf_file.read()
+    pdf_stream = io.BytesIO(pdf_data)
+    response = requests.post(
+        'https://pdftotext.com/ExtractText',
+        files={'pdffile': pdf_stream},
+        data={'form': 'pdftotext'}
+    )
+    text = response.text.strip()
+
+    # Tokenize question and text
+    input_ids = tokenizer.encode(question, text)
+
+    # Perform question answering
+    outputs = model(torch.tensor([input_ids]), return_dict=True)
+    answer_start = outputs.start_logits.argmax().item()
+    answer_end = outputs.end_logits.argmax().item()
+    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(input_ids[answer_start:answer_end+1]))
+
+    return answer
+
+inputs = [
+    gr.inputs.File(label="PDF document"),
+    gr.inputs.Textbox(label="Question")
+]
+
+outputs = gr.outputs.Textbox(label="Answer")
+
+gr.Interface(fn=answer_question, inputs=inputs, outputs=outputs, title="PDF Question Answering Tool", 
+             description="Upload a PDF document and ask a question. The app will use a pre-trained model to find the answer.").launch()