import streamlit as st
from hf_model import load_model, summarize_text, answer_question

# Title and Sidebar
st.title("Legal NLP Application")
st.sidebar.title("Options")

# Select the task
task = st.sidebar.radio("Select Task", ["Summarization", "Question Answering"])

# Task 1: Summarization
if task == "Summarization":
    st.header("Summarize Legal Text")
    user_input = st.text_area("Enter Legal Text", height=200, placeholder="Paste a legal text or case summary here...")

    if st.button("Summarize"):
        st.write("Loading model...")
        model = load_model(task="summarization")  # Load summarization model
        summary = summarize_text(model, user_input)
        st.write("### Summary")
        st.write(summary)

# Task 2: Question Answering
elif task == "Question Answering":
    st.header("Legal Question Answering")
    context = st.text_area("Enter Context", height=200, placeholder="Paste a legal document or case description here...")
    question = st.text_input("Ask a Legal Question", placeholder="Type your legal question here...")

    if st.button("Get Answer"):
        st.write("Loading model...")
        model = load_model(task="question-answering")  # Load QA model
        answer = answer_question(model, question, context)
        st.write("### Answer")
        st.write(answer)