import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from sentence_transformers import SentenceTransformer
import pickle

st.title("Fin$mart Chatbot")
st.markdown("Ask financial questions and get answers based on expert knowledge.")

# Load models
@st.cache_resource
def load_models():
    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
    embedder = SentenceTransformer("all-MiniLM-L6-v2")
    return tokenizer, model, embedder

tokenizer, model, embedder = load_models()

# Load vector store from existing pickle file
@st.cache_resource
def load_vector_store():
    with open("vectorstore.pkl", "rb") as f:
        index, texts, _ = pickle.load(f)  # We ignore embeddings if not needed
    return index, texts

index, texts = load_vector_store()

# Chat interface
prompt = st.chat_input("Ask something about finance...")

if prompt:
    # Embed query and retrieve top 3 results
    q_embed = embedder.encode([prompt])
    _, I = index.search(q_embed, k=3)
    context = " ".join([texts[i] for i in I[0]])

    # Build input for Flan-T5
    input_text = (
        f"You are a helpful financial assistant. Use the information provided below to answer the user's question.\n\n"
        f"Context: {context}\n\n"
        f"Question: {prompt}\n\n"
        f"Answer:"
    )

    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
    outputs = model.generate(**inputs, max_length=150)
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Display response
    st.markdown(f"**Answer:** {answer}")

    # Show retrieved context
    with st.expander("Context Used"):
        for i in I[0]:
            st.write(texts[i])