app.py

import streamlit as st
import torch
from transformers import T5TokenizerFast, AutoModel


DEVICE = "cpu"

@st.cache_resource()
def load_model():
    model = AutoModel.from_pretrained("ohgnues/FiDT5", trust_remote_code=True).to(DEVICE)
    tokenizer = T5TokenizerFast.from_pretrained("ohgnues/FiDT5")
    
    return model, tokenizer

def main():
    st.title("Question Answering System")

    with st.expander("Context 1"):
        context1 = st.text_area("Enter Context 1", value="Does He Love You \"Does He Love You\" is a song written by Sandy Knox and Billy Stritch, and recorded as a duet by American country music artists Reba McEntire and Linda Davis. It was released in August 1993 as the first single from Reba's album \"Greatest Hits Volume Two\". It is one of country music's several songs about a love triangle. \"Does He Love You\" was written in 1982 by Billy Stritch. He recorded it with a trio in which he performed at the time, because he wanted a song that could be sung by the other two members")

    with st.expander("Context 2"):
        context2 = st.text_area("Enter Context 2")

    with st.expander("Context 3"):
        context3 = st.text_area("Enter Context 3")

    question = st.text_input("Enter Question", value="who sings does he love me with reba")

    if st.button("Show Answer"):
        model, tokenizer = load_model()
        answer = get_answer(model, tokenizer, context1, context2, context3, question)
        st.write("Answer:", answer)

def get_answer(model, tokenizer, context1, context2, context3, question):
    inputs = [question + " " + context for context in [context1, context2, context3] if context]
    input_tokens = tokenizer(inputs, return_tensors="pt", padding=True).to(DEVICE)

    output_tokens = model.generate(input_ids=input_tokens.input_ids, attention_mask=input_tokens.attention_mask, max_length=16)

    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)

if __name__ == "__main__":
    main()