Spaces:
Runtime error
Runtime error
File size: 2,652 Bytes
bc12901 bc6a638 bc12901 bc6a638 8171e8e bc6a638 8171e8e bc12901 bc6a638 2919076 bc6a638 bc12901 bc6a638 bc12901 bc6a638 8171e8e bc12901 8171e8e bc12901 8171e8e bc12901 |
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 |
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, load_document
def ensure_list(x):
if isinstance(x, list):
return x
else:
return [x]
@st.experimental_singleton
def construct_pipeline():
device = "cuda" if torch.cuda.is_available() else "cpu"
ret = get_pipeline(device=device)
return ret
@st.cache
def run_pipeline(question, document):
return construct_pipeline()(question=question, **document.context)
st.markdown("# DocQuery: Query Documents w/ NLP")
if "document" not in st.session_state:
st.session_state["document"] = None
input_type = st.radio("Pick an input type", ["Upload", "URL"], horizontal=True)
def load_file_cb():
if st.session_state.file_input is None:
return
file = st.session_state.file_input
with loading_placeholder:
with st.spinner("Processing..."):
document = load_bytes(file, file.name)
_ = document.context
st.session_state.document = document
def load_url_cb():
if st.session_state.url_input is None:
return
url = st.session_state.url_input
with loading_placeholder:
with st.spinner("Downloading..."):
document = load_document(url)
with st.spinner("Processing..."):
_ = document.context
st.session_state.document = document
if input_type == "Upload":
file = st.file_uploader(
"Upload a PDF or Image document", key="file_input", on_change=load_file_cb
)
elif input_type == "URL":
# url = st.text_input("URL", "", on_change=load_url_callback, key="url_input")
url = st.text_input("URL", "", key="url_input", on_change=load_url_cb)
question = st.text_input("QUESTION", "")
document = st.session_state.document
loading_placeholder = st.empty()
if document is not None:
col1, col2 = st.columns(2)
col1.image(document.preview, use_column_width=True)
if document is not None and question is not None and len(question) > 0:
predictions = run_pipeline(question=question, document=document)
col2.header("Answers")
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)"
|