import streamlit as st

# Import transformer classes for generaiton
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, GPT2Tokenizer, GPT2LMHeadModel, GPT2Model
# Import torch for datatype attributes
import torch
# Import the prompt wrapper...but for llama index
from llama_index.prompts.prompts import SimpleInputPrompt
# Import the llama index HF Wrapper
from llama_index.llms import HuggingFaceLLM
# Bring in embeddings wrapper
from llama_index.embeddings import LangchainEmbedding
# Bring in HF embeddings - need these to represent document chunks
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
# Bring in stuff to change service context
from llama_index import set_global_service_context
from llama_index import ServiceContext
# Import deps to load documents
from llama_index import VectorStoreIndex, SimpleDirectoryReader
from pathlib import Path
import pypdf
import time
import os
import tempfile

# Define variable to hold llama2 weights namingfiner
name = "gpt2"
# Set auth token variable from hugging face
auth_token = "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"

@st.cache_resource
def get_tokenizer_model():
    # Create tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained(name, use_auth_token=auth_token)

    # Create model
    model = GPT2LMHeadModel.from_pretrained(name)
    
    return model, tokenizer

model, tokenizer = get_tokenizer_model()

# disclaimer
st.title('LLM Deployment Prototype for Production')
st.caption("Special thanks to my mentor, Medkham Chanthavong, for all the support in making this project a reality. Connect with Medkham on LinkedIn: [Medkham Chanthavong](https://www.linkedin.com/in/medkham-chanthavong-4524b0125/).")

st.subheader('Disclaimer')
st.write("Due to Streamlit's lack of cloud-GPU support and limited memory, we're using GPT-2 instead of Llama 2 7B. GPT-2, \
being smaller, is less capable of accurately retrieving information from external data. This site is a prototype for LLM deployment \
and may not retrieve precise data from external sources. Additionally, the site might experience crashes due to memory constraints.\
In case of a crash, please feel free shoot me a message at my [LinkedIn profile](https://www.linkedin.com/in/duwe-ng/) so I can reboot the site. \
For a more advanced model demonstration using Llama 2 7B, check out our Colab notebook: \
[LLM Deployment Prototype](https://colab.research.google.com/drive/1bGf9rKntMjH4KtpKs9ryucj1nbiKs_zk?usp=sharing).")


# # Initialize the SimpleInputPrompt with an empty template
# query_wrapper_prompt = SimpleInputPrompt("{query_str}")

# # Streamlit UI to let the user update the system prompt
# # Start with an empty string or a default prompt
# default_prompt = ""
# user_system_prompt = st.text_area("How can I best assist you?", value="", height=100)
# update_button = st.button('Request')


# Import the prompt wrapper for llama index
# Create a system prompt 
system_prompt = """
Hey there! 👋 You are here to help you out in the best way you can. Think of yourself as your friendly and trustworthy assistant. Your top priority is to be super helpful and safe in your responses. \
Here's a quick heads up about what you can expect from me: \
1. Positive Vibes Only: I steer clear of anything harmful, unethical, or offensive. No racism, sexism, or toxic stuff here. I'm all about being respectful and keeping things on the up-and-up. \
2. Making Sense Matters: If the user question seems a bit confusing or doesn't quite add up, You will let you know and try to clarify things instead of giving the user a misleading answer. \
3. Honesty is Key: Not sure about something? You won't make stuff up. If you don't have the answer, you will be upfront about it. \
4. All About Your PDFs: The user documents are your focus. Got a question about the content in your PDFs? That's exactly what you are here for. Let's dive into those documents and find the answers the user need! \
So, how can you assist the user today with their PDFs?
"""

# Throw together the query wrapper
query_wrapper_prompt = SimpleInputPrompt("{query_str}")

# Initialize the llm object with a placeholder or default system prompt
llm = HuggingFaceLLM(
    context_window=1024,
    max_new_tokens=128,
    system_prompt="",
    query_wrapper_prompt=query_wrapper_prompt,
    model=model,
    tokenizer=tokenizer
)

# # Function to update the system prompt and reinitialize the LLM with the new prompt
# def update_system_prompt(new_prompt):
#     global llm
#     llm.system_prompt = new_prompt


# if update_button:
#     # Update the system prompt and reinitialize the LLM
#     update_system_prompt(user_system_prompt)
#     st.success('Requested')
    
# check if user request save to the memory
# st.write(llm.system_prompt)

# Create and dl embeddings instance
embeddings=LangchainEmbedding(
    HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
)

# Create new service context instance
service_context = ServiceContext.from_defaults(
    chunk_size=256,
    llm=llm,
    embed_model=embeddings
)

# And set the service context
set_global_service_context(service_context)


# Upload PDF and process it
st.subheader("Please upload your data to enable the retrieval system; otherwise, it will not respond.")

documents = []
uploaded_file = st.file_uploader("Upload PDF", type="pdf")

if uploaded_file is not None:
    # Create a temporary directory
    with tempfile.TemporaryDirectory() as temp_dir:
        # Write the uploaded file to a file in the temporary directory
        temp_file_path = os.path.join(temp_dir, uploaded_file.name)
        with open(temp_file_path, "wb") as f:
            f.write(uploaded_file.getbuffer())

        # Now you can use SimpleDirectoryReader on the temp_dir
        documents = SimpleDirectoryReader(temp_dir).load_data()
        st.success(f"File {uploaded_file.name} uploaded successfully.")
        
index = VectorStoreIndex.from_documents(documents)
        
# Setup index query engine using LLM
query_engine = index.as_query_engine(streaming=True, similarity_top_k=1)

# Create centered main title
st.title('👔 HireMind 🧩')
        
# setup a session to hold all the old prompt
if 'messages' not in st.session_state:
    st.session_state.messages = []
        
# print out the history message
for message in st.session_state.messages:
    st.chat_message(message['role']).markdown(message['content'])
        
        
# Create a text input box for the user
# If the user hits enter
prompt = st.chat_input('Input your prompt here')

if prompt:
    st.chat_message('user').markdown(prompt)
    st.session_state.messages.append({'role': 'user', 'content': prompt})
    
    response = query_engine.query(prompt)
    
    st.chat_message('assistant').markdown(response)
    st.session_state.messages.append(
        {'role': 'assistant', 'content': response}
    )


langchain.embeddings.huggingface.HuggingFaceEmbeddings
[((1422, 1473), 'streamlit.title', 'st.title', (['"""LLM Deployment Prototype for Production"""'], {}), "('LLM Deployment Prototype for Production')\n", (1430, 1473), True, 'import streamlit as st\n'), ((1474, 1718), 'streamlit.caption', 'st.caption', (['"""Special thanks to my mentor, Medkham Chanthavong, for all the support in making this project a reality. Connect with Medkham on LinkedIn: [Medkham Chanthavong](https://www.linkedin.com/in/medkham-chanthavong-4524b0125/)."""'], {}), "(\n    'Special thanks to my mentor, Medkham Chanthavong, for all the support in making this project a reality. Connect with Medkham on LinkedIn: [Medkham Chanthavong](https://www.linkedin.com/in/medkham-chanthavong-4524b0125/).'\n    )\n", (1484, 1718), True, 'import streamlit as st\n'), ((1710, 1736), 'streamlit.subheader', 'st.subheader', (['"""Disclaimer"""'], {}), "('Disclaimer')\n", (1722, 1736), True, 'import streamlit as st\n'), ((1737, 2482), 'streamlit.write', 'st.write', (['"""Due to Streamlit\'s lack of cloud-GPU support and limited memory, we\'re using GPT-2 instead of Llama 2 7B. GPT-2, being smaller, is less capable of accurately retrieving information from external data. This site is a prototype for LLM deployment and may not retrieve precise data from external sources. Additionally, the site might experience crashes due to memory constraints.In case of a crash, please feel free shoot me a message at my [LinkedIn profile](https://www.linkedin.com/in/duwe-ng/) so I can reboot the site. For a more advanced model demonstration using Llama 2 7B, check out our Colab notebook: [LLM Deployment Prototype](https://colab.research.google.com/drive/1bGf9rKntMjH4KtpKs9ryucj1nbiKs_zk?usp=sharing)."""'], {}), '(\n    "Due to Streamlit\'s lack of cloud-GPU support and limited memory, we\'re using GPT-2 instead of Llama 2 7B. GPT-2, being smaller, is less capable of accurately retrieving information from external data. This site is a prototype for LLM deployment and may not retrieve precise data from external sources. Additionally, the site might experience crashes due to memory constraints.In case of a crash, please feel free shoot me a message at my [LinkedIn profile](https://www.linkedin.com/in/duwe-ng/) so I can reboot the site. For a more advanced model demonstration using Llama 2 7B, check out our Colab notebook: [LLM Deployment Prototype](https://colab.research.google.com/drive/1bGf9rKntMjH4KtpKs9ryucj1nbiKs_zk?usp=sharing)."\n    )\n', (1745, 2482), True, 'import streamlit as st\n'), ((4064, 4096), 'llama_index.prompts.prompts.SimpleInputPrompt', 'SimpleInputPrompt', (['"""{query_str}"""'], {}), "('{query_str}')\n", (4081, 4096), False, 'from llama_index.prompts.prompts import SimpleInputPrompt\n'), ((4176, 4335), 'llama_index.llms.HuggingFaceLLM', 'HuggingFaceLLM', ([], {'context_window': '(1024)', 'max_new_tokens': '(128)', 'system_prompt': '""""""', 'query_wrapper_prompt': 'query_wrapper_prompt', 'model': 'model', 'tokenizer': 'tokenizer'}), "(context_window=1024, max_new_tokens=128, system_prompt='',\n    query_wrapper_prompt=query_wrapper_prompt, model=model, tokenizer=tokenizer\n    )\n", (4190, 4335), False, 'from llama_index.llms import HuggingFaceLLM\n'), ((4953, 5030), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'chunk_size': '(256)', 'llm': 'llm', 'embed_model': 'embeddings'}), '(chunk_size=256, llm=llm, embed_model=embeddings)\n', (4981, 5030), False, 'from llama_index import ServiceContext\n'), ((5076, 5119), 'llama_index.set_global_service_context', 'set_global_service_context', (['service_context'], {}), '(service_context)\n', (5102, 5119), False, 'from llama_index import set_global_service_context\n'), ((5150, 5263), 'streamlit.subheader', 'st.subheader', (['"""Please upload your data to enable the retrieval system; otherwise, it will not respond."""'], {}), "(\n    'Please upload your data to enable the retrieval system; otherwise, it will not respond.'\n    )\n", (5162, 5263), True, 'import streamlit as st\n'), ((5286, 5328), 'streamlit.file_uploader', 'st.file_uploader', (['"""Upload PDF"""'], {'type': '"""pdf"""'}), "('Upload PDF', type='pdf')\n", (5302, 5328), True, 'import streamlit as st\n'), ((5897, 5939), 'llama_index.VectorStoreIndex.from_documents', 'VectorStoreIndex.from_documents', (['documents'], {}), '(documents)\n', (5928, 5939), False, 'from llama_index import VectorStoreIndex, SimpleDirectoryReader\n'), ((6089, 6113), 'streamlit.title', 'st.title', (['"""👔 HireMind 🧩"""'], {}), "('👔 HireMind 🧩')\n", (6097, 6113), True, 'import streamlit as st\n'), ((6482, 6521), 'streamlit.chat_input', 'st.chat_input', (['"""Input your prompt here"""'], {}), "('Input your prompt here')\n", (6495, 6521), True, 'import streamlit as st\n'), ((1200, 1262), 'transformers.GPT2Tokenizer.from_pretrained', 'GPT2Tokenizer.from_pretrained', (['name'], {'use_auth_token': 'auth_token'}), '(name, use_auth_token=auth_token)\n', (1229, 1262), False, 'from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, GPT2Tokenizer, GPT2LMHeadModel, GPT2Model\n'), ((1295, 1332), 'transformers.GPT2LMHeadModel.from_pretrained', 'GPT2LMHeadModel.from_pretrained', (['name'], {}), '(name)\n', (1326, 1332), False, 'from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, GPT2Tokenizer, GPT2LMHeadModel, GPT2Model\n'), ((4841, 4893), 'langchain.embeddings.huggingface.HuggingFaceEmbeddings', 'HuggingFaceEmbeddings', ([], {'model_name': '"""all-MiniLM-L6-v2"""'}), "(model_name='all-MiniLM-L6-v2')\n", (4862, 4893), False, 'from langchain.embeddings.huggingface import HuggingFaceEmbeddings\n'), ((6583, 6652), 'streamlit.session_state.messages.append', 'st.session_state.messages.append', (["{'role': 'user', 'content': prompt}"], {}), "({'role': 'user', 'content': prompt})\n", (6615, 6652), True, 'import streamlit as st\n'), ((6761, 6837), 'streamlit.session_state.messages.append', 'st.session_state.messages.append', (["{'role': 'assistant', 'content': response}"], {}), "({'role': 'assistant', 'content': response})\n", (6793, 6837), True, 'import streamlit as st\n'), ((5404, 5433), 'tempfile.TemporaryDirectory', 'tempfile.TemporaryDirectory', ([], {}), '()\n', (5431, 5433), False, 'import tempfile\n'), ((5543, 5585), 'os.path.join', 'os.path.join', (['temp_dir', 'uploaded_file.name'], {}), '(temp_dir, uploaded_file.name)\n', (5555, 5585), False, 'import os\n'), ((5816, 5879), 'streamlit.success', 'st.success', (['f"""File {uploaded_file.name} uploaded successfully."""'], {}), "(f'File {uploaded_file.name} uploaded successfully.')\n", (5826, 5879), True, 'import streamlit as st\n'), ((6329, 6361), 'streamlit.chat_message', 'st.chat_message', (["message['role']"], {}), "(message['role'])\n", (6344, 6361), True, 'import streamlit as st\n'), ((6538, 6561), 'streamlit.chat_message', 'st.chat_message', (['"""user"""'], {}), "('user')\n", (6553, 6561), True, 'import streamlit as st\n'), ((6709, 6737), 'streamlit.chat_message', 'st.chat_message', (['"""assistant"""'], {}), "('assistant')\n", (6724, 6737), True, 'import streamlit as st\n'), ((5764, 5795), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', (['temp_dir'], {}), '(temp_dir)\n', (5785, 5795), False, 'from llama_index import VectorStoreIndex, SimpleDirectoryReader\n')]