import streamlit as st
import torch
from transformers import AutoProcessor, UdopForConditionalGeneration
# from datasets import load_dataset

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

processor = AutoProcessor.from_pretrained("microsoft/udop-large", apply_ocr=True)
model = UdopForConditionalGeneration.from_pretrained("microsoft/udop-large")

question = "Question answering. How many unsafe practice of Lifting Operation"

st.title("CIC Demo (by ITT)")
st.write("Upload and Select a document (/an image) to test the model.")

# File selection
uploaded_files = st.file_uploader("Upload document(s) [/image(s)]:", type=["docx", "pdf", "pptx", "jpg", "jpeg", "png"], accept_multiple_files=True)
selected_file = st.selectbox("Select a document (/an image):", uploaded_files, format_func=lambda file: file.name if file else "None")

# Display selected file
if selected_file is not None and selected_file != "None":
    file_extension = selected_file.name.split(".")[-1]
    if file_extension in ["jpg", "jpeg", "png"]:
        st.image(selected_file, caption="Selected Image")
    else:
        st.write("Selected file: ", selected_file.name)

# Model testing button
testButton = st.button("Test Model")
if testButton and selected_file != "None":
    st.write("Testing the model with the selected image...")
    encoding = processor(image, question, words, boxes=boxes, return_tensors="pt")
    predicted_ids = model.generate(**encoding)
    print(processor.batch_decode(predicted_ids, skip_special_tokens=True)[0])    
elif testButton and selected_file == "None":
    st.write("Please upload and select a document (/an image).")