# app.py
import streamlit as st
from embedding import setup_embeddings, load_data
from model_utils import load_llama_model, is_finance_question, ask_finance_bot

# Load data and embeddings
questions, answers = load_data()
embedding_model, index = setup_embeddings(answers)
tokenizer, model = load_llama_model()

# Streamlit UI
st.set_page_config(page_title="DiMowkayBot - Finance Assistant", layout="centered")
st.title("DiMowkayBot - Your Finance Q&A Assistant")

user_query = st.text_input("Ask a finance-related question:")

if user_query:
    with st.spinner("Thinking..."):
        if not is_finance_question(user_query, tokenizer, model):
            st.warning("I'm specialized in finance and can't help with that. How can I assist you with a finance-related question today?")
        else:
            answer = ask_finance_bot(user_query, answers, embedding_model, index, tokenizer, model)
            st.success("Response:")
            st.write(answer)