import streamlit as st
import os
from llama_index import VectorStoreIndex, ServiceContext
from llama_index.llms import HuggingFaceLLM
from llama_index.prompts.prompts import SimpleInputPrompt
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from llama_index.embeddings import LangchainEmbedding

# Function to set up the Llama2 model and index
def setup_llama2():
    system_prompt = """
    You are a Q&A assistant. Your goal is to answer questions as
    accurately as possible based on the instructions and context provided.
    """

    query_wrapper_prompt = SimpleInputPrompt("{query_str}")

    llm = HuggingFaceLLM(
        context_window=4096,
        max_new_tokens=256,
        generate_kwargs={"temperature": 0.0, "do_sample": False},
        system_prompt=system_prompt,
        query_wrapper_prompt=query_wrapper_prompt,
        tokenizer_name="meta-llama/Llama-2-7b-chat-hf",
        model_name="meta-llama/Llama-2-7b-chat-hf",
        device_map="auto",
        model_kwargs={"torch_dtype": torch.float16, "load_in_8bit": True}
    )

    embed_model = LangchainEmbedding(
        HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
    )

    service_context = ServiceContext.from_defaults(
        chunk_size=1024,
        llm=llm,
        embed_model=embed_model
    )

    index = VectorStoreIndex(service_context=service_context)

    return index

# Streamlit app
def main():
    st.title("Q&A Assistant with Llama2 and Hugging Face")

    # Upload PDF file
    pdf_file = st.file_uploader("Upload a PDF file", type=["pdf"])

    if pdf_file is not None:
        # Process the uploaded PDF file (you may need to adapt this based on your needs)
        # For example, you might want to convert the PDF to text or extract relevant information.
        # Here, we assume the content is stored in a variable named 'pdf_content'.
        pdf_content = pdf_file.read()

        # Set up the Llama2 model and index
        index = setup_llama2()

        # Section for asking questions
        st.header("Ask a Question")

        # Query input
        query = st.text_input("Enter your question:")

        if st.button("Submit Question"):
            if query:
                # Execute the query
                response = index.as_query_engine().query(query)

                # Display the response
                st.header("Answer")
                st.markdown(response)

# Run the Streamlit app
if __name__ == "__main__":
    main()