app.py

import os
import time
import base64
import logging
import torch
import streamlit as st
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.llms import HuggingFacePipeline
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain.embeddings import HuggingFaceBgeEmbeddings 
from langchain.llms import HuggingFacePipeline
from langchain.vectorstores import Chroma



@st.cache_resource(show_spinner=False)
def load_model(model_name):
    logger.info("Loading model ..")
    start_time = time.time()

    if model_name=='llama':
        from langchain.llms import CTransformers

        model =  CTransformers(model="TheBloke/Llama-2-7B-Chat-GGML", model_file = 'llama-2-7b-chat.ggmlv3.q2_K.bin', 
                    model_type='llama', gpu_layers=0, config={"context_length":2048,})
        tokenizer = None
    
    elif model_name=='mistral':
        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

        model_id="filipealmeida/Mistral-7B-Instruct-v0.1-sharded"

        quant_config = BitsAndBytesConfig(
           load_in_4bit=True,
           bnb_4bit_quant_type="nf4",
           bnb_4bit_use_double_quant=True,
           bnb_4bit_compute_dtype=torch.bfloat16)
        
        model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, quantization_config=quant_config, device_map="auto")

        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        tokenizer.pad_token = tokenizer.eos_token
    
    logger.info(f"Model Loading Time : {time.time() - start_time} .")

    return model, tokenizer


@st.cache_resource(show_spinner=False)
def load_db(device, local_embed=False,  CHROMA_PATH = './ChromaDB'):
    """
    Load vector embeddings and Chroma database 
    """
    encode_kwargs = {'normalize_embeddings': True}
    embed_id = "BAAI/bge-large-en-v1.5"    
    start_time = time.time()

    #TODO : LOOK INTO LOADING ONLY A SINGLE FILE FROM HF REPO TO REDUCE MEMORY
    if local_embed: 
        from transformers import AutoModel

        PATH_TO_EMBEDDING_FOLDER = ""
        # TODO : load only pytorch bin file
        embeddings = AutoModel.from_pretrained(PATH_TO_EMBEDDING_FOLDER, trust_remote_code=True)
        embeddings = HuggingFaceBgeEmbeddings(model_name="  ", model_kwargs={"trust_remote_code":True}) 
        logger.info('Loading embeddings locally.')  
 
   
    else:
        embeddings = HuggingFaceBgeEmbeddings(model_name=embed_id , model_kwargs={"device": device}, encode_kwargs=encode_kwargs)
        logger.info('Loading embeddings from Hub.')
        

    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings)
    logger.info(f"Vector Embeddings and Chroma Database Loading Time : {time.time() - start_time} .")
    return db


def wrap_model(model, tokenizer):
    """wrap transformers pipeline with HuggingFacePipeline
    """
    text_generation_pipeline = pipeline(
        model=model,
        tokenizer=tokenizer,
        task="text-generation",
        temperature=0.2,
        repetition_penalty=1.1,
        #return_full_text=True,
        max_new_tokens=1000,
        pad_token_id=2,
        do_sample=True)
    HF_pipeline = HuggingFacePipeline(pipeline=text_generation_pipeline)
    return HF_pipeline



def fetch_context(db, model, model_name, query, template, use_compressor=True):
    """
    Perform similarity search and retrieve related context to query.
    I have stored large documents in db so I can apply compressor on the set of retrived documents to 
    make sure that returned compressed context is relevant to the query.
    """
    if use_compressor:
        if model_name=='llama':
            compressor = LLMChainExtractor.from_llm(model)
            compressor.llm_chain.prompt.template = template['llama_rag_template']
            
        elif model_name=='mistral':
            global HF_pipeline_model
            HF_pipeline_model = wrap_model(model)
            compressor = LLMChainExtractor.from_llm(HF_pipeline_model)
            compressor.llm_chain.prompt.template = template['rag_template']
        
        retriever = db.as_retriever(search_type = "mmr") 
        compression_retriever = ContextualCompressionRetriever(base_compressor=compressor,
                                                        base_retriever=retriever)
        #logger.info(f"User Query : {query}")
        compressed_docs = compression_retriever.get_relevant_documents(query)
        #logger.info(f"Retrieved Compressed Docs : {compressed_docs}")

        return compressed_docs
    
    docs = db.max_marginal_relevance_search(query)
    #logger.info(f"Retrieved Docs : {docs}")

    return docs


def format_context(docs):
    """
    clean and format chunks into documents to pass as context
    """
    cleaned_docs = [doc for doc in docs if ">>>" not in doc.page_content]
    return "\n\n".join(doc.page_content for doc in cleaned_docs)



def llm_chain_with_context(model, model_name, query, context, template):
    """ 
    Run simple chain with formatted prompt including query and retrieved context and the underlying model to generate a response.
    """
    formated_context = format_context(context)
    # Give a precise answer to the question based on the context. Don't be verbose.
    if model_name=='llama':
        prompt_template = PromptTemplate(input_variables=['context', 'user_query'], template = template['llama_prompt_template'])
        llm_chain = LLMChain(llm=model, prompt=prompt_template) 
        
    elif model_name=='mistral':
        prompt_template = PromptTemplate(input_variables=['context', 'user_query'], template = template['prompt_template'])
        llm_chain = LLMChain(llm=HF_pipeline_model, prompt=prompt_template) 

    output = llm_chain.predict(user_query=query, context=formated_context)
    return output


def generate_response(query,  model, template):
    start_time = time.time()
    progress_text = "Running Inference. Please wait."
    my_bar = st.progress(0, text=progress_text)
    # fill those as appropriate
    my_bar.progress(0.3, "Loading Model.  Please wait.")
    time.sleep(1) 
    my_bar.progress(0.5, "Running RAG.  Please wait.")
    context = fetch_context(db, model, model_name, query, template)

    my_bar.progress(0.7, "Generating Answer.  Please wait.")
    response = llm_chain_with_context(model, model_name, query, context, template)

    logger.info(f"Total Execution Time: {time.time() - start_time}") 

    my_bar.progress(0.9, "Post Processing.  Please wait.")
    
    my_bar.progress(1.0, "Done")
    time.sleep(1) 
    my_bar.empty()  
    return response


# show background image
def convert_to_base64(bin_file):
    with open(bin_file, 'rb') as f:
        data = f.read()
    return base64.b64encode(data).decode()

def set_as_background_img(png_file):
    bin_str = convert_to_base64(png_file) 
    background_img = '''
    <link href='https://fonts.googleapis.com/css?family=Libre Baskerville' rel='stylesheet'>
    <style>
    .stApp {
    background-image: url("data:image/png;base64,%s");
    background-size: cover;
    background-repeat: no-repeat;
    background-attachment: scroll; 
    }
    </style>
     ''' % bin_str
    st.markdown(background_img, unsafe_allow_html=True)
    return   


def stream_to_screen(response):
    print(response)
    for word in response.split():
        yield word + " "
        time.sleep(0.15)
    
    
if __name__=="__main__":

    st.set_page_config(page_title='StoicCyber', page_icon="🏛️", layout="centered", initial_sidebar_state="collapsed")
    set_as_background_img('pxfuel.jpg')
    # header
    original_title = '<h1 style="font-family: Libre Baskerville; color:#faf8f8; font-size: 30px; text-align: left; ">STOIC Ω CYBER</h1>'
    st.markdown(original_title, unsafe_allow_html=True)
    
    
    # hide footer and header  
    hide_st_style = """
                <style>
                header {visibility: hidden;}
                footer {visibility: hidden;}
                </style>
                """
    st.markdown(hide_st_style, unsafe_allow_html=True)
    
    logger = logging.getLogger(__name__)
    logging.basicConfig(
        filename="app.log",
        filemode="a",
        format="%(asctime)s.%(msecs)03d %(levelname)s [%(funcName)s] %(message)s",
        level=logging.INFO,
        datefmt="%Y-%m-%d %H:%M:%S",)


    device = "cuda" if torch.cuda.is_available() else "cpu"
    model_name = "llama" if device=="cpu" else "mistral"
    logger.info(f"Running {model_name} model for inference on {device}")

    all_templates = { "llama_prompt_template" : """<s>[INST]\n<<SYS>>\nYou are a stoic teacher that provide guidance and advice inspired by Stoic philosophy on navigating life's challenges with resilience and inner peace. Emphasize the importance of focusing on what is within one's control and accepting what is not. Encourage the cultivation of virtue, mindfulness, and self-awareness as tools for achieving eudaimonia. Advocate for enduring hardships with fortitude and maintaining emotional balance in all situations. Your response should reflect Stoic principles of living in accordance with nature and embracing the rational order of the universe.
            You should guide the reader towards a fulfilling life focused on virtue rather than external things because living in accordance with virtue leads to eudaimonia or flourishing.\n\n
            Give a precise answer to the question based on the context. Don't be verbose.\n\n
            context:
            {context}\n<</SYS>>\n\n
            question:
            {user_query}
            [/INST]""",

            "llama_rag_template" :"""<s>[INST]\n<<SYS>>\nGiven the following question and context, summarize the parts that are relevant to answer the question. If none of the context is relevant return NO_OUTPUT.\n\n>
            - Do not mention quotes.\n\n
            - Reply using a single sentence.\n\n
            > Context:\n
            >>>\n{context}\n>>>\n<</SYS>>\n\n
            Question: {question}\n
            [/INST]
            The relevant parts of the context are:
            """,

            "prompt_template":"""You are a stoic teacher that provide guidance and advice inspired by Stoic philosophy on navigating life's challenges with resilience and inner peace. Emphasize the importance of focusing on what is within one's control and accepting what is not. Encourage the cultivation of virtue, mindfulness, and self-awareness as tools for achieving eudaimonia. Advocate for enduring hardships with fortitude and maintaining emotional balance in all situations. Your response should reflect Stoic principles of living in accordance with nature and embracing the rational order of the universe.
                You should guide the reader towards a fulfilling life focused on virtue rather than external things because living in accordance with virtue leads to eudaimonia or flourishing.
                context:
                {context}

                question:
                {user_query}

                Answer:
                """,
            "rag_prompt" : """Given the following question and context, summarize the parts that are relevant to answer the question. If none of the context is relevant return NO_OUTPUT.\n\n>
                - Do not mention quotes.\n\n>
                - Reply using a single sentence.\n\n>

                Question: {question}\n> Context:\n>>>\n{context}\n>>>\nRelevant parts"""}

    # Loading and caching db and model
    #my_bar = st.progress(0, "Loading Database.  Please wait.")
    #my_bar.progress(0.1, "Loading Embedding & Database.  Please wait.")
    #my_bar.progress(0.7, "Loading Model.  Please wait.")

    db = load_db(device)
    model, tokenizer = load_model(model_name)
    
    response = """Ah, a question that touches the very heart of Stoic philosophy! Mindfulness, my dear, is an interior discipline that allows you to align your thoughts and actions with the rational order of the universe. To cultivate mindfulness, one must practice self-awareness, for it is only by understanding thyself that thou canst understand the world around thee. 

        Here are a few practical steps to cultivate mindfulness in accordance with Stoic principles:
        
             1. Observe your thoughts and emotions without judgment. As the ancient Stoic philosopher Epictetus once said, "The things within our control, are those over which we have the power of decision; and those things are external to us, over which we have no control." Recognize that which lies within thy control and accept what is not. 
            
             2. Practice meditation and inner reflection. Meditation allows thee to quiet the mind and listen to the whispers of thy inner being. Observe thy thoughts and emotions with detachment, as if observing a passing cloud. Through this practice, thou shalt become attuned to the ebb and flow of life, ever-aware of the present moment.
        """
    
    user_question = st.chat_input('What do you want to ask ..')

    
    # streamlit chat
    if user_question is not None and user_question!="":
        with st.chat_message("Human", avatar="🧔🏻"):
            st.write(user_question)

        #response = generate_response(user_question,  model, all_templates)
        response = ""
        with st.chat_message("AI", avatar="🏛️"):
            st.write_stream(stream_to_screen(response))