import streamlit as st
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.vectorstores import FAISS
from transformers import pipeline

import sentence_transformers
print(sentence_transformers.__version__)
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

import subprocess
import sys

# Install sentence-transformers if not installed
try:
    import sentence_transformers
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "sentence-transformers"])


# Initialize embedding model
embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")

def chunk_text(text, chunk_size=500):
    words = text.split()
    chunks = [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
    return chunks

# Streamlit app
st.title("Simple RAG Application")
data = st.text_area("Paste your text here:")
if data:
    text_chunks = chunk_text(data)
    vectorstore = FAISS.from_texts(text_chunks, embeddings)
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    
    question = st.text_input("Ask a question:")
    if question:
        relevant_docs = retriever.get_relevant_documents(question)
        context = " ".join([doc.page_content for doc in relevant_docs])
        answer = qa_pipeline(question=question, context=context)
        st.write("Answer:", answer["answer"])