import os

os.environ["TOKENIZERS_PARALLELISM"] = "false"

import streamlit as st

import torch
from docquery.pipeline import get_pipeline
from docquery.document import load_bytes

device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = get_pipeline(device=device)


def process_document(file, question):
    # prepare encoder inputs
    document = load_document(file.name)
    return pipeline(question=question, **document.context)


def ensure_list(x):
    if isinstance(x, list):
        return x
    else:
        return [x]


st.title("DocQuery: Query Documents Using NLP")
file = st.file_uploader("Upload a PDF or Image document")
question = st.text_input("QUESTION", "")

document = None

if file is not None:
    col1, col2 = st.columns(2)

    document = load_bytes(file, file.name)
    col1.image(document.preview, use_column_width=True)

if document is not None and question is not None and len(question) > 0:
    predictions = pipeline(question=question, **document.context)

    col2.header("Probabilities")
    for p in ensure_list(predictions):
        col2.subheader(f"{ p['answer'] }: { round(p['score'] * 100, 1)}%")


"DocQuery uses LayoutLMv1 fine-tuned on DocVQA, a document visual question answering dataset, as well as SQuAD, which boosts its English-language comprehension. To use it, simply upload an image or PDF, type a question, and click 'submit', or click one of the examples to load them."

"[Github Repo](https://github.com/impira/docquery)"