Greatmonkey's picture
Create app.py
bacc64d verified
import streamlit as st
from transformers import AutoModelForQuestionAnswering, AutoTokenizer
def main():
# Custom CSS for styling
custom_css = """
<style>
body {
font-family: 'Arial', sans-serif;
background-color: #FF0000; /* Set your desired background color here */
}
.title {
font-size: 2.5em;
color: #333;
text-align: center;
padding: 1em;
background-color: #3498db;
color: #fff;
border-radius: 10px;
margin-bottom: 1em;
}
.input-container {
margin: 2em;
}
.button-container {
text-align: center;
}
.result-container {
margin: 2em;
padding: 1em;
background-color: #fff;
border-radius: 10px;
box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
}
.answer {
background-color: #3498db;
padding: 10px;
border-radius: 5px;
margin-top: 10px;
}
img {
max-width: 100%;
height: auto;
border-radius: 10px;
box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
}
</style>
"""
st.markdown(custom_css, unsafe_allow_html=True)
# Title
st.markdown("<div class='title'>Question Answering with Transformers</div>", unsafe_allow_html=True)
# Model Selection Dropdown
model_name = st.selectbox("Select Model", ["deepset/roberta-base-squad2", "bert-large-uncased-whole-word-masking-finetuned-squad", "distilbert-base-cased-distilled-squad",
"bert-base-uncased",
"albert-base-v2"])
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Get user input
st.markdown("<div class='input-container'>", unsafe_allow_html=True)
context = st.text_area("Enter the context (max 400 words):")
question = st.text_input("Enter your question:")
st.markdown("</div>", unsafe_allow_html=True)
if st.button("Get Answer"):
#st.markdown("<div class='result-container'>", unsafe_allow_html=True)
if not question or not context:
st.warning("Please enter both a question and a context.")
else:
# Tokenize input
try:
# Check word count in the context
if len(context.split()) > 400:
raise ValueError("Context exceeds 400 words limit.")
inputs = tokenizer.encode_plus(question, context, return_tensors='pt')
# Get predictions
outputs = model(**inputs)
start_logits = outputs.start_logits
end_logits = outputs.end_logits
# Get top N answer spans
top_n = 3
start_indexes = start_logits.argsort(dim=1, descending=True)[:, :top_n]
end_indexes = end_logits.argsort(dim=1, descending=True)[:, :top_n]
# Display detailed answers
st.subheader(f"Question: {question}")
for i in range(top_n):
start_index = start_indexes[0, i].item()
end_index = end_indexes[0, i].item()
answer = tokenizer.decode(inputs['input_ids'][0, start_index:end_index + 1])
# Highlight answer in context
highlighted_context = f"{context[:start_index]}**{context[start_index:end_index+1]}**{context[end_index+1:]}"
# Display confidence scores
confidence_start = start_logits[0, start_index].item()
confidence_end = end_logits[0, end_index].item()
if answer == "":
continue
else:
st.markdown(f"<div class='answer'><strong>Answer:</strong> {answer}<br>"
f"<strong>Confidence (Start):</strong> {confidence_start:.4f}<br>"
f"<strong>Confidence (End):</strong> {confidence_end:.4f}</div>", unsafe_allow_html=True)
except ValueError as ve:
st.error(str(ve))
except Exception as e:
st.error(f"An error occurred: {e}")
st.markdown("</div>", unsafe_allow_html=True)
if __name__ == "__main__":
main()