import streamlit as st
from model import load_model, generate_answer,find_context
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from dataset import dataset_a
from sentence_transformers import SentenceTransformer

# Load model automatically
MODEL_PATH = "namngo/CDS_vit5"
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH)
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
pos_sentences = dataset_a["context"]
model2=SentenceTransformer('namngo/CDS_retrival')

# Streamlit UI customization
st.set_page_config(page_title="Chat Bot Công dân số", page_icon="🤖", layout="wide")
st.title("🤖 Chat Bot Công dân số")
st.markdown("---")
st.success("✅ Model Loaded Successfully")

st.sidebar.header("⚙️ Settings")
max_length = st.sidebar.slider("Max Answer Length", min_value=50, max_value=500, value=256, step=10)

st.subheader("📌 Ask a Question")
# context = st.text_area("📝 Context:", height=150)
question = st.text_input("❓ Question:")
context=find_context(pos_sentences,question,model2)


if st.button("🚀 Generate Answer"):
    answer = generate_answer(model, tokenizer, context, question, max_length=max_length)
    st.markdown("### 💡 Answer:")
    st.info(answer)