import gradio as gr
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, LayoutLMv3ImageProcessor

model_name = "TusharGoel/LiLT-Document-QA"

revision = "3a510b84c579386c5edfd3881ba839bba28e6a44"
tokenizer = AutoTokenizer.from_pretrained(model_name, apply_ocr = True, revision=revision)
image_processor = LayoutLMv3ImageProcessor()

model = AutoModelForQuestionAnswering.from_pretrained(model_name, revision=revision)
model.eval()

def qna(image, question):
    try:
        res = image_processor(image, apply_ocr = True)
        words = res["words"][0]
        boxes = res["boxes"][0]
    
        encoding = tokenizer(question, words, boxes = boxes, return_token_type_ids=True, return_tensors="pt", truncation=True, padding="max_length")
    
        word_ids = encoding.word_ids(0)
        outputs = model(**encoding)
    
        start_scores = outputs.start_logits
        end_scores = outputs.end_logits
    
        start, end = word_ids[start_scores.argmax(-1).item()], word_ids[end_scores.argmax(-1).item()]
    
        answer = " ".join(words[start : end + 1])

    except:
        answer = "No Answer"
    
    
    return answer


img = gr.Image(label="Image")
question = gr.Text(label="Question")
label = gr.Label(label="label")

iface = gr.Interface(fn=qna, inputs=[img, question], outputs=label, title="LiLT - Document Question Answering", allow_duplication=True)
iface.launch()