import streamlit as st
import pickle
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import pipeline

# Load saved quotes and embeddings from root directory
with open("quote_embeddings.pkl", "rb") as f:
    quotes, embeddings = pickle.load(f)

# Initialize embedder and FAISS index
embedder = SentenceTransformer('all-MiniLM-L6-v2')

embeddings = embeddings.astype('float32')
index = faiss.IndexFlatIP(embeddings.shape[1])
faiss.normalize_L2(embeddings)
index.add(embeddings)

# Initialize text generator pipeline
generator = pipeline('text-generation', model='distilgpt2')

# Define RAG search function
def rag_search(query, top_k=3):
    q_emb = embedder.encode([query]).astype('float32')
    faiss.normalize_L2(q_emb)
    scores, indices = index.search(q_emb, top_k)
    context = "\n".join([f"{quotes[i]['quote']} — {quotes[i].get('author','Unknown')}" for i in indices[0]])
    prompt = f"Answer using these quotes:\n{context}\nQuestion: {query}\nAnswer:"
    outputs = generator(prompt, max_length=100, num_return_sequences=1)
    answer = outputs[0]['generated_text'].split('Answer:')[-1].strip()
    return answer

# Streamlit UI starts here
st.title("RAG Quote-Based Q&A")

user_query = st.text_input("Ask a question related to quotes:")

if user_query:
    with st.spinner("Generating answer..."):
        answer = rag_search(user_query)
    st.markdown("### Answer:")
    st.write(answer)