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
def load_model(model_name, logger, ):
    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
def load_db(logger, 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="whatever-model-you-are-using", model_kwargs={"trust_remote_code":True}) 
        logger.info('Loading embeddings locally.')  
        # Test the local embeddings
        embed = embeddings.get_text_embedding("Hello World!")
        print(len(embed))
        print(embed[:5])
   
    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, query, logger, 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':
            HF_pipeline_model = wrap_model(model)
            global HF_pipeline_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, logger):
    """ 
    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, logger):
    start_time = time.time()
    progress_text = "Loading model. Please wait."
    my_bar = st.progress(0, text=progress_text)
    context = fetch_context(db, model, model_name, query, template, logger)
    # fill those as appropriate
    my_bar.progress(0.1, "Loading Database.  Please wait.")

    my_bar.progress(0.3, "Loading Model.  Please wait.")
    
    my_bar.progress(0.5, "Running RAG.  Please wait.")

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

    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   

    
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)
    
    user_question = st.chat_input('What do you want to ask ..')
    
    # hide footer and header  
    hide_st_style = """
                <style>
                header {visibility: hidden;}
                footer {visibility: hidden;}
                </style>
                """
    st.markdown(hide_st_style, unsafe_allow_html=True)
    
    # set logger
    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",)


    # model to use in spaces depends on the available device 
    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.
            context:
            {context}\n<</SYS>>\n\n
            question:
            {user_query}
            [/INST]""",

            "llmaa_rag_prompt" :"""<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"""}


    db = load_db(logger, device)

    model, tokenizer = load_model(model_name, logger)
    
    # 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, logger)   
        with st.chat_message("AI", avatar="🏛️"):
            st.write(response)