SPINIX/app.py at main · Rishi-Raj-Anand/SPINIX · GitHub

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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import streamlit as st
from transformers import pipeline

st.set_page_config(layout="wide", page_title="SPINIX")

# Cache model
@st.cache_resource
def load_qa_model():
    return pipeline(
        "question-answering",
        model="distilbert-base-cased-distilled-squad" # default to 30522 tokens
    )

# loading model and caching
qa_pipeline = load_qa_model()

#  initializating session states
if "submitted_context" not in st.session_state:
    st.session_state.submitted_context = ""

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []


# Title
st.markdown("<h1 style='text-align: center;'>SPINIX</h1>", unsafe_allow_html=True)

# 2 columns
col_left, col_right = st.columns([1, 1], gap="small")


# left- context part
with col_left:
    st.markdown("<h3 style='text-align: center;'>Context</h3>", unsafe_allow_html=True)

    current_context_input = st.text_area(
        "Paste your text below:",
        height=400,
        placeholder="Paste your article, document, or text here...",
        label_visibility="collapsed",
        key="context_input_area"
    )

    # Submit Button
    if st.button("Submit Context", width="stretch"):
        if current_context_input.strip():
            st.session_state.submitted_context = current_context_input

            # clear context on submission of new context
            st.session_state.chat_history = []
            st.toast("Context submitted successfully! You can now ask questions.")
        else:
            st.toast("Please enter some text first")


#  right - chat
with col_right:
    st.markdown("<h3 style='text-align: center;'>Ask Questions</h3>", unsafe_allow_html=True)

    chat_container = st.container(height=400, border=True)

    # Display Chat History
    with chat_container:

        for chat in st.session_state.chat_history:
            with st.chat_message("user"):
                st.write(chat['question'])

            with st.chat_message("assistant"):
                st.write(chat['answer'])

    prompt = st.chat_input("Enter your question...", key="chat_input_widget")

    if prompt:

        if not st.session_state.submitted_context:
             st.toast("Please submit a context first")
        else:

            with chat_container:
                 with st.chat_message("user"):
                     st.write(prompt)

                 # Process answer
                 with st.chat_message("assistant"):
                     with st.spinner("Thinking..."):
                         # Run model using session state context
                         result = qa_pipeline(
                             question=prompt,
                             context=st.session_state.submitted_context
                         )

                         answer_text = ""

                         # threshold
                         if result['score'] > 0.1:
                             answer_text = result['answer']
                         else:
                             answer_text = "I cannot find the answer to that question in the provided context."

                         st.write(answer_text)

            # add new question and answer
            st.session_state.chat_history.append(
                {'question': prompt, 'answer': answer_text}
            )

            # Rerun to update new messages
            st.rerun()