import streamlit as st
from transformers import pipeline
import pandas as pd
import re

# Load the Question Answering model
qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")

# Load SOP Dataset
@st.cache
def load_sop_dataset():
    """Load SOP dataset from CSV."""
    dataset = pd.read_csv("dataset.csv")  # Ensure this file is uploaded to your Hugging Face Space
    return dataset

# Load the dataset
dataset = load_sop_dataset()

# Utility function to find relevant contexts
def find_relevant_contexts(question, dataset):
    """Search for relevant contexts in the dataset."""
    relevant_contexts = []
    for index, row in dataset.iterrows():
        if re.search(question, row["text"], re.IGNORECASE):
            relevant_contexts.append(row["text"])
    return relevant_contexts

# Streamlit UI
st.title("SOP Question Answering AI")
st.markdown("Ask any question about Standard Operating Procedures:")

# User input
question = st.text_area("Enter your question:", "")
specific_context = st.checkbox("Use specific SOP context?")

context = None
if specific_context:
    st.write("Choose a context:")
    context = st.selectbox("SOP Contexts", dataset["text"])
else:
    if question:
        st.write("Searching for relevant contexts...")
        relevant_contexts = find_relevant_contexts(question, dataset)
        if relevant_contexts:
            context = st.selectbox("Relevant SOP Contexts", relevant_contexts)
        else:
            st.warning("No relevant contexts found. Try refining your question.")

# Generate answer
if st.button("Get Answer"):
    if context:
        with st.spinner("Finding the answer..."):
            result = qa_pipeline(question=question, context=context)
            st.success("Answer:")
            st.write(result["answer"])
            st.write("Confidence Score:", result["score"])
    else:
        st.warning("Please select a context or refine your question.")