import streamlit as st
import torch
from transformers import (
    AutoTokenizer, AutoModelForCausalLM,
    BitsAndBytesConfig,
    TextStreamer,
)
import whisper
import os

############ config ############
# general config
whisper_model_names=["tiny", "base", "small", "medium", "large"]
data_root_path = os.path.join('.','data')
file_types = ['pdf','png','jpg','wav']
for filetype in file_types:
    if not os.path.exists(os.path.join(data_root_path,filetype)):
        os.makedirs(os.path.join(data_root_path,filetype))

# streamlit config
## Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "Type a message to start a conversation"}]





############ User Interface ############
# Title
st.title('LLAMA RAG Demo')
st.divider()

st.title('Model name and auth token')
# Configs
model_name = st.text_input('Enter your Hugging Face model name', value="meta-llama/Llama-2-7b-chat-hf")
auth_token = st.text_input('Enter your Hugging Face auth token', value="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")
system_prompt = st.text_area('Enter your system prompt', value="You are a helpful, respectful and honest assistant.")
whisper_model_name = st.selectbox('Select your whisper model',options=whisper_model_names)
use_cuda = st.checkbox('Use CUDA', value=True)
isfile = False
## File uploader
from streamlit import file_uploader
uploadedfile = file_uploader("Choose a \"PDF\" file (now support only pdf)")
if uploadedfile is not None:
    isfile = True
    with open(os.path.join(data_root_path,'pdf',uploadedfile.name),"wb") as f:
        f.write(uploadedfile.getbuffer())
    st.success("File uploaded successfully : {}".format(uploadedfile.name))
st.divider()





############ function ############
def clear_chat_history():
    st.session_state.messages = [{"role": "assistant", "content": "Type a message to start a conversation"}]
    for file in os.listdir(os.path.join(data_root_path,'pdf')):
        os.remove(os.path.join(data_root_path,'pdf',file))
    
st.button('Clear Chat History', on_click=clear_chat_history)

# Load Tokenizer and Model
@st.cache_resource
def get_tokenizer_model():
    # Create tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir='./model/', token=auth_token)
    # Create model
    quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_compute_dtype=torch.float16,
    )
    model = AutoModelForCausalLM.from_pretrained(model_name,
            cache_dir='./model/', token=auth_token,
            quantization_config=quantization_config,
            # rope_scaling={"type":"dynamic", "factor":2},
            max_memory=f'{int(torch.cuda.mem_get_info()[0]/1024**3)-2}GB'
    )
    return tokenizer, model

# RAG engine
def get_rag_queryengine(_tokenizer, model, system_prompt):
    from llama_index.prompts.prompts import SimpleInputPrompt
    from llama_index.llms import HuggingFaceLLM
    system_prompt_ = f"[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
    query_wrapper_prompt = SimpleInputPrompt("{query_str} [/INST]")
    llm = HuggingFaceLLM(context_window=4096,
                    max_new_tokens=256,
                    system_prompt=system_prompt_,
                    query_wrapper_prompt=query_wrapper_prompt,
                    model=model,
                    tokenizer=_tokenizer
                    )
    
    # Create embeddings
    from llama_index.embeddings import LangchainEmbedding
    from langchain.embeddings.huggingface import HuggingFaceEmbeddings
    embeddings=LangchainEmbedding(
    HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    )
    from llama_index import ServiceContext
    from llama_index import set_global_service_context
    service_context = ServiceContext.from_defaults(
        chunk_size=1024,
        llm=llm,
        embed_model=embeddings
    )
    set_global_service_context(service_context)

    from llama_index import VectorStoreIndex, download_loader
    PyMuPDFReader = download_loader("PyMuPDFReader")
    loader = PyMuPDFReader()
    for file in os.listdir(os.path.join(data_root_path,'pdf')):
        # !!! This is not a good way to load data. I will fix this later
        # this makes the only last file in the folder to be loaded
        documents = loader.load_data(file_path=os.path.join(data_root_path,'pdf',file), metadata=True)
    index = VectorStoreIndex.from_documents(documents)
    query_engine = index.as_query_engine()
    return query_engine

# whisper
def whisper_stt(*,model, device, audio_path)->str:
    # load model
    # # model : model name of whisper. default is base
    # # devie : argument from args. default is cpu
    audio_model = whisper.load_model(model,device)
    # stt - audio.wav
    result = audio_model.transcribe(audio_path)
    # return result : str list
    return result["text"]





############ main ############
# Load Tokenizer and Model, RAG engine
tokenizer, model = get_tokenizer_model()
if isfile:
    engine = get_rag_queryengine(tokenizer, model, system_prompt)

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

prompt = st.chat_input('User: ')
if prompt:
    # update(append) chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.write(prompt)

# Here... text streamer does not work as well as I intended with streamlit
# I will try to fix this later
if st.session_state.messages[-1]["role"] == "user":
    if isfile:
        with st.chat_message("assistant"):
            # model inference
            output_text = engine.query(prompt)

            placeholder = st.empty()
            placeholder.markdown(output_text)
        st.session_state.messages.append({"role": "assistant", "content": output_text})

    else:
        with st.chat_message("assistant"):
            # model inference
            output_text = "Please upload a file first"
            placeholder = st.empty()
            placeholder.markdown(output_text)
        st.session_state.messages.append({"role": "assistant", "content": output_text})


langchain.embeddings.huggingface.HuggingFaceEmbeddings
[((306, 331), 'os.path.join', 'os.path.join', (['"""."""', '"""data"""'], {}), "('.', 'data')\n", (318, 331), False, 'import os\n'), ((772, 798), 'streamlit.title', 'st.title', (['"""LLAMA RAG Demo"""'], {}), "('LLAMA RAG Demo')\n", (780, 798), True, 'import streamlit as st\n'), ((799, 811), 'streamlit.divider', 'st.divider', ([], {}), '()\n', (809, 811), True, 'import streamlit as st\n'), ((813, 850), 'streamlit.title', 'st.title', (['"""Model name and auth token"""'], {}), "('Model name and auth token')\n", (821, 850), True, 'import streamlit as st\n'), ((874, 969), 'streamlit.text_input', 'st.text_input', (['"""Enter your Hugging Face model name"""'], {'value': '"""meta-llama/Llama-2-7b-chat-hf"""'}), "('Enter your Hugging Face model name', value=\n    'meta-llama/Llama-2-7b-chat-hf')\n", (887, 969), True, 'import streamlit as st\n'), ((978, 1081), 'streamlit.text_input', 'st.text_input', (['"""Enter your Hugging Face auth token"""'], {'value': '"""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"""'}), "('Enter your Hugging Face auth token', value=\n    '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')\n", (991, 1081), True, 'import streamlit as st\n'), ((1093, 1199), 'streamlit.text_area', 'st.text_area', (['"""Enter your system prompt"""'], {'value': '"""You are a helpful, respectful and honest assistant."""'}), "('Enter your system prompt', value=\n    'You are a helpful, respectful and honest assistant.')\n", (1105, 1199), True, 'import streamlit as st\n'), ((1216, 1286), 'streamlit.selectbox', 'st.selectbox', (['"""Select your whisper model"""'], {'options': 'whisper_model_names'}), "('Select your whisper model', options=whisper_model_names)\n", (1228, 1286), True, 'import streamlit as st\n'), ((1297, 1332), 'streamlit.checkbox', 'st.checkbox', (['"""Use CUDA"""'], {'value': '(True)'}), "('Use CUDA', value=True)\n", (1308, 1332), True, 'import streamlit as st\n'), ((1416, 1475), 'streamlit.file_uploader', 'file_uploader', (['"""Choose a "PDF" file (now support only pdf)"""'], {}), '(\'Choose a "PDF" file (now support only pdf)\')\n', (1429, 1475), False, 'from streamlit import file_uploader\n'), ((1722, 1734), 'streamlit.divider', 'st.divider', ([], {}), '()\n', (1732, 1734), True, 'import streamlit as st\n'), ((2038, 2098), 'streamlit.button', 'st.button', (['"""Clear Chat History"""'], {'on_click': 'clear_chat_history'}), "('Clear Chat History', on_click=clear_chat_history)\n", (2047, 2098), True, 'import streamlit as st\n'), ((5259, 5282), 'streamlit.chat_input', 'st.chat_input', (['"""User: """'], {}), "('User: ')\n", (5272, 5282), True, 'import streamlit as st\n'), ((2212, 2298), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['model_name'], {'cache_dir': '"""./model/"""', 'token': 'auth_token'}), "(model_name, cache_dir='./model/', token=\n    auth_token)\n", (2241, 2298), False, 'from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextStreamer\n'), ((2339, 2445), 'transformers.BitsAndBytesConfig', 'BitsAndBytesConfig', ([], {'load_in_4bit': '(True)', 'bnb_4bit_quant_type': '"""nf4"""', 'bnb_4bit_compute_dtype': 'torch.float16'}), "(load_in_4bit=True, bnb_4bit_quant_type='nf4',\n    bnb_4bit_compute_dtype=torch.float16)\n", (2357, 2445), False, 'from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextStreamer\n'), ((3074, 3114), 'llama_index.prompts.prompts.SimpleInputPrompt', 'SimpleInputPrompt', (['"""{query_str} [/INST]"""'], {}), "('{query_str} [/INST]')\n", (3091, 3114), False, 'from llama_index.prompts.prompts import SimpleInputPrompt\n'), ((3125, 3297), 'llama_index.llms.HuggingFaceLLM', 'HuggingFaceLLM', ([], {'context_window': '(4096)', 'max_new_tokens': '(256)', 'system_prompt': 'system_prompt_', 'query_wrapper_prompt': 'query_wrapper_prompt', 'model': 'model', 'tokenizer': '_tokenizer'}), '(context_window=4096, max_new_tokens=256, system_prompt=\n    system_prompt_, query_wrapper_prompt=query_wrapper_prompt, model=model,\n    tokenizer=_tokenizer)\n', (3139, 3297), False, 'from llama_index.llms import HuggingFaceLLM\n'), ((3786, 3864), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'chunk_size': '(1024)', 'llm': 'llm', 'embed_model': 'embeddings'}), '(chunk_size=1024, llm=llm, embed_model=embeddings)\n', (3814, 3864), False, 'from llama_index import ServiceContext\n'), ((3899, 3942), 'llama_index.set_global_service_context', 'set_global_service_context', (['service_context'], {}), '(service_context)\n', (3925, 3942), False, 'from llama_index import set_global_service_context\n'), ((4026, 4058), 'llama_index.download_loader', 'download_loader', (['"""PyMuPDFReader"""'], {}), "('PyMuPDFReader')\n", (4041, 4058), False, 'from llama_index import VectorStoreIndex, download_loader\n'), ((4407, 4449), 'llama_index.VectorStoreIndex.from_documents', 'VectorStoreIndex.from_documents', (['documents'], {}), '(documents)\n', (4438, 4449), False, 'from llama_index import VectorStoreIndex, download_loader\n'), ((4720, 4753), 'whisper.load_model', 'whisper.load_model', (['model', 'device'], {}), '(model, device)\n', (4738, 4753), False, 'import whisper\n'), ((5332, 5401), 'streamlit.session_state.messages.append', 'st.session_state.messages.append', (["{'role': 'user', 'content': prompt}"], {}), "({'role': 'user', 'content': prompt})\n", (5364, 5401), True, 'import streamlit as st\n'), ((1937, 1972), 'os.path.join', 'os.path.join', (['data_root_path', '"""pdf"""'], {}), "(data_root_path, 'pdf')\n", (1949, 1972), False, 'import os\n'), ((3607, 3659), 'langchain.embeddings.huggingface.HuggingFaceEmbeddings', 'HuggingFaceEmbeddings', ([], {'model_name': '"""all-MiniLM-L6-v2"""'}), "(model_name='all-MiniLM-L6-v2')\n", (3628, 3659), False, 'from langchain.embeddings.huggingface import HuggingFaceEmbeddings\n'), ((4115, 4150), 'os.path.join', 'os.path.join', (['data_root_path', '"""pdf"""'], {}), "(data_root_path, 'pdf')\n", (4127, 4150), False, 'import os\n'), ((5175, 5207), 'streamlit.chat_message', 'st.chat_message', (["message['role']"], {}), "(message['role'])\n", (5190, 5207), True, 'import streamlit as st\n'), ((5217, 5248), 'streamlit.markdown', 'st.markdown', (["message['content']"], {}), "(message['content'])\n", (5228, 5248), True, 'import streamlit as st\n'), ((5411, 5434), 'streamlit.chat_message', 'st.chat_message', (['"""user"""'], {}), "('user')\n", (5426, 5434), True, 'import streamlit as st\n'), ((5444, 5460), 'streamlit.write', 'st.write', (['prompt'], {}), '(prompt)\n', (5452, 5460), True, 'import streamlit as st\n'), ((5847, 5926), 'streamlit.session_state.messages.append', 'st.session_state.messages.append', (["{'role': 'assistant', 'content': output_text}"], {}), "({'role': 'assistant', 'content': output_text})\n", (5879, 5926), True, 'import streamlit as st\n'), ((6157, 6236), 'streamlit.session_state.messages.append', 'st.session_state.messages.append', (["{'role': 'assistant', 'content': output_text}"], {}), "({'role': 'assistant', 'content': output_text})\n", (6189, 6236), True, 'import streamlit as st\n'), ((424, 462), 'os.path.join', 'os.path.join', (['data_root_path', 'filetype'], {}), '(data_root_path, filetype)\n', (436, 462), False, 'import os\n'), ((484, 522), 'os.path.join', 'os.path.join', (['data_root_path', 'filetype'], {}), '(data_root_path, filetype)\n', (496, 522), False, 'import os\n'), ((1539, 1593), 'os.path.join', 'os.path.join', (['data_root_path', '"""pdf"""', 'uploadedfile.name'], {}), "(data_root_path, 'pdf', uploadedfile.name)\n", (1551, 1593), False, 'import os\n'), ((1992, 2033), 'os.path.join', 'os.path.join', (['data_root_path', '"""pdf"""', 'file'], {}), "(data_root_path, 'pdf', file)\n", (2004, 2033), False, 'import os\n'), ((5648, 5676), 'streamlit.chat_message', 'st.chat_message', (['"""assistant"""'], {}), "('assistant')\n", (5663, 5676), True, 'import streamlit as st\n'), ((5782, 5792), 'streamlit.empty', 'st.empty', ([], {}), '()\n', (5790, 5792), True, 'import streamlit as st\n'), ((5951, 5979), 'streamlit.chat_message', 'st.chat_message', (['"""assistant"""'], {}), "('assistant')\n", (5966, 5979), True, 'import streamlit as st\n'), ((6092, 6102), 'streamlit.empty', 'st.empty', ([], {}), '()\n', (6100, 6102), True, 'import streamlit as st\n'), ((4339, 4380), 'os.path.join', 'os.path.join', (['data_root_path', '"""pdf"""', 'file'], {}), "(data_root_path, 'pdf', file)\n", (4351, 4380), False, 'import os\n'), ((2716, 2741), 'torch.cuda.mem_get_info', 'torch.cuda.mem_get_info', ([], {}), '()\n', (2739, 2741), False, 'import torch\n')]