app.py

import os

# Set the cache directory to persistent storage
os.environ["HF_HOME"] = "/data/.cache/huggingface"

import streamlit as st
from utils.help import get_disclaimer
from utils.format import sec_to_time, fix_latex, get_youtube_embed
from utils.rag_utils import load_youtube_data, load_book_data, load_summary, fixed_knn_retrieval, get_random_question
from utils.system_prompts import get_expert_system_prompt, get_synthesis_system_prompt
from utils.openai_utils import embed_question_openai, openai_domain_specific_answer_generation, openai_context_integration
from utils.llama_utils import get_bnb_config, load_base_model, load_fine_tuned_model, generate_response

st.set_page_config(page_title="AI University")

st.markdown("""
    <style>
    .video-wrapper {
        position: relative;
        padding-bottom: 56.25%;
        height: 0;
    }
    .video-wrapper iframe {
        position: absolute;
        top: 0;
        left: 0;
        width: 100%;
        height: 100%;
    }
    </style>
    """, unsafe_allow_html=True)

# ---------------------------------------
# paths
# ---------------------------------------
base_path = "data/"
base_model_path = "meta-llama/Llama-3.2-11B-Vision-Instruct"
base_model_path_3B = "meta-llama/Llama-3.2-3B-Instruct"
adapter_path = "./LLaMA-TOMMI-1.0/"

st.title(":red[AI University] :gray[/] FEM")
# st.markdown("### Finite Element Method")
st.markdown("Welcome to :red[AI University]—an AI-powered platform designed to address scientific course queries, dynamically adapting to instructors' teaching styles and students' learning needs. This prototype demonstrates the capabilities of the AI University platform by providing expert answers to queries related to a graduate-level :red[Finite Element Method (FEM)] course")

st.markdown(" ")
st.markdown(" ")
# st.divider()
# Sidebar for settings
with st.sidebar:
    st.header("Settings")
    # with st.container(border=True):
    # Embedding model

    model_name = st.selectbox("Choose content embedding model", [
        "text-embedding-3-small",
        # "text-embedding-3-large",
        # "all-MiniLM-L6-v2", 
        # "all-mpnet-base-v2"
    ], 
    # help="""
    # Select the embedding model to use for encoding the retrieved text data. 
    # Options include OpenAI's `text-embedding-3` models and two widely
    # used SentenceTransformers models.
    # """
    )
    
    with st.container(border=True):
        st.write('**Video lectures**')
        yt_token_choice = st.select_slider("Token per content", [256, 512, 1024], value=256, help="Larger values lead to an increase in the length of each retrieved piece of content", key="yt_token_len")
        yt_chunk_tokens = yt_token_choice
        yt_max_content = {128: 32, 256: 16, 512: 8, 1024: 4}[yt_chunk_tokens]
        top_k_YT = st.slider("Number of relevant content pieces to retrieve", 0, yt_max_content, 4, key="yt_token_num")
        yt_overlap_tokens = yt_chunk_tokens // 4

    # st.divider()
    with st.container(border=True):
        st.write('**Textbook**')
        show_textbook = False
        # show_textbook = st.toggle("Show Textbook Content", value=False)
        latex_token_choice = st.select_slider("Token per content", [128, 256, 512, 1024], value=256, help="Larger values lead to an increase in the length of each retrieved piece of content", key="latex_token_len")
        latex_chunk_tokens = latex_token_choice
        latex_max_content = {128: 32, 256: 16, 512: 8, 1024: 4}[latex_chunk_tokens]
        top_k_Latex = st.slider("Number of relevant content pieces to retrieve", 0, latex_max_content, 4, key="latex_token_num")
        # latex_overlap_tokens = latex_chunk_tokens // 4
        latex_overlap_tokens = 0

    st.write(' ')
    with st.expander('Expert model', expanded=False):
            use_expert_answer = st.toggle("Use expert answer", value=True)
            show_expert_responce = st.toggle("Show initial expert answer", value=False)

            st.session_state.expert_model = st.selectbox(
                "Choose the LLM model",
                ["LLaMA-TOMMI-1.0-11B", "LLaMA-3.2-11B", "gpt-4o-mini"],
                index=0,
                key='a1model'
            )

            if st.session_state.expert_model in ["LLaMA-TOMMI-1.0-11B", "LLaMA-3.2-11B"]:
                expert_do_sample = st.toggle("Enable Sampling", value=False, key='expert_sample')

                if expert_do_sample:
                    expert_temperature = st.slider("Temperature", 0.0, 1.5, 0.7, key='expert_temp')
                    expert_top_k = st.slider("Top K", 0, 100, 50, key='expert_top_k')
                    expert_top_p = st.slider("Top P", 0.0, 1.0, 0.95, key='expert_top_p')
                else:
                    expert_num_beams = st.slider("Num Beams", 1, 4, 1, key='expert_num_beams')

                expert_max_new_tokens = st.slider("Max New Tokens", 100, 2000, 500, step=50, key='expert_max_new_tokens')
            else:
                expert_api_temperature = st.slider("Temperature", 0.0, 1.5, 0.7, key='a1t')
                expert_api_top_p = st.slider("Top P", 0.0, 1.0, 0.9, key='a1p')

    with st.expander('Synthesis model',expanded=False):
        # with st.container(border=True):
        # Choose the LLM model
        st.session_state.synthesis_model = st.selectbox(
            "Choose the LLM model",
            ["LLaMA-3.2-3B","gpt-4o-mini"], # "LLaMA-3.2-11B", 
            index=1,
            key='a2model'
        )
        
        if st.session_state.synthesis_model in ["LLaMA-3.2-3B", "LLaMA-3.2-11B"]:
                synthesis_do_sample = st.toggle("Enable Sampling", value=False, key='synthesis_sample')

                if synthesis_do_sample:
                    synthesis_temperature = st.slider("Temperature", 0.0, 1.5, 0.7, key='synthesis_temp')
                    synthesis_top_k = st.slider("Top K", 0, 100, 50, key='synthesis_top_k')
                    synthesis_top_p = st.slider("Top P", 0.0, 1.0, 0.95, key='synthesis_top_p')
                else:
                    synthesis_num_beams = st.slider("Num Beams", 1, 4, 1, key='synthesis_num_beams')

                synthesis_max_new_tokens = st.slider("Max New Tokens", 100, 2000, 1500, step=50, key='synthesis_max_new_tokens')
        else:
            # Temperature
            synthesis_api_temperature = st.slider("Temperature", 0.0, .3, .5, help="Defines the randomness in the next token prediction. Lower: More predictable and focused. Higher: More adventurous and diverse.", key='a2t')

            synthesis_api_top_p = st.slider("Top P", 0.1, 0.5, .3, help="Defines the range of token choices the model can consider in the next prediction. Lower: More focused and restricted to high-probability options. Higher: More creative, allowing consideration of less likely options.", key='a2p')

# Main content area
if "question" not in st.session_state:
    st.session_state.question = ""


text_area_placeholder = st.empty()
question_help = "Including details or instructions improves the answer."
st.session_state.question = text_area_placeholder.text_area(
    "**Enter your query about Finite Element Method**",
    height=120,
    value=st.session_state.question,
    help=question_help
)

_, col1, col2, _ = st.columns([4, 2, 4, 3]) 
with col1:
    submit_button_placeholder = st.empty()

with col2:
    if st.button("Random Question"):
        while True:
            random_question = get_random_question(base_path + "/questions.txt")
            if random_question != st.session_state.question:
                break
        st.session_state.question = random_question
        text_area_placeholder.text_area(
            "**Enter your query about Finite Element Method:**",
            height=120,
            value=st.session_state.question,
            help=question_help
        )

with st.spinner("Loading LLaMA-TOMMI-1.0-11B..."):
    if st.session_state.expert_model == "LLaMA-TOMMI-1.0-11B":
        if 'tommi_model' not in st.session_state: 
            tommi_model, tommi_tokenizer = load_fine_tuned_model(adapter_path, base_model_path)
            st.session_state.tommi_model = tommi_model
            st.session_state.tommi_tokenizer = tommi_tokenizer


with st.spinner("Loading LLaMA-3.2-11B..."):
    if "LLaMA-3.2-11B" in [st.session_state.expert_model, st.session_state.synthesis_model]:
        if 'llama_model' not in st.session_state: 
            llama_model, llama_tokenizer = load_base_model(base_model_path)
            st.session_state.llama_model = llama_model
            st.session_state.llama_tokenizer = llama_tokenizer

with st.spinner("Loading LLaMA-3.2-3B..."):
    if "LLaMA-3.2-3B" in [st.session_state.expert_model, st.session_state.synthesis_model]:
        if 'llama_model_3B' not in st.session_state: 
            llama_model_3B, llama_tokenizer_3B = load_base_model(base_model_path_3B)
            st.session_state.llama_model_3B = llama_model_3B
            st.session_state.llama_tokenizer_3B = llama_tokenizer_3B

# Load YouTube and LaTeX data
text_data_YT, context_embeddings_YT = load_youtube_data(base_path, model_name, yt_chunk_tokens, yt_overlap_tokens)
text_data_Latex, context_embeddings_Latex = load_book_data(base_path, model_name, latex_chunk_tokens, latex_overlap_tokens)
summary = load_summary('data/KG_FEM_summary.json')

if 'question_answered' not in st.session_state:
    st.session_state.question_answered = False
if 'context_by_video' not in st.session_state:
    st.session_state.context_by_video = {}
if 'context_by_section' not in st.session_state:
    st.session_state.context_by_section = {}
if 'answer' not in st.session_state:
    st.session_state.answer = ""
if 'playing_video_id' not in st.session_state:
    st.session_state.playing_video_id = None

if submit_button_placeholder.button("AI Answer", type="primary"): 
    if st.session_state.question == "":
        st.markdown("")
        st.write("Please enter a query. :smirk:")
        st.session_state.question_answered = False

    else:    
        with st.spinner("Finding relevant contexts..."):
            question_embedding = embed_question_openai(st.session_state.question, model_name)
            initial_max_k = int(0.1 * context_embeddings_YT.shape[0])
            idx_YT = fixed_knn_retrieval(question_embedding, context_embeddings_YT, top_k=top_k_YT, min_k=0)
            idx_Latex = fixed_knn_retrieval(question_embedding, context_embeddings_Latex, top_k=top_k_Latex, min_k=0)

        with st.spinner("Answering the question..."):
            relevant_contexts_YT = sorted([text_data_YT[i] for i in idx_YT], key=lambda x: x['order'])
            relevant_contexts_Latex = sorted([text_data_Latex[i] for i in idx_Latex], key=lambda x: x['order'])

            st.session_state.context_by_video = {}
            for context_item in relevant_contexts_YT:
                video_id = context_item['video_id']
                if video_id not in st.session_state.context_by_video:
                    st.session_state.context_by_video[video_id] = []
                st.session_state.context_by_video[video_id].append(context_item)

            st.session_state.context_by_section = {}
            for context_item in relevant_contexts_Latex:
                section_id = context_item['section']
                if section_id not in st.session_state.context_by_section:
                    st.session_state.context_by_section[section_id] = []
                st.session_state.context_by_section[section_id].append(context_item)

            context = ''
            for i, (video_id, contexts) in enumerate(st.session_state.context_by_video.items(), start=1):
                for context_item in contexts:
                    start_time = int(context_item['start'])
                    context += f'Video {i}, time: {sec_to_time(start_time)}:' + context_item['text'] + '\n\n'

            for i, (section_id, contexts) in enumerate(st.session_state.context_by_section.items(), start=1):
                context += f'Section {i} ({section_id}):\n'
                for context_item in contexts:
                    context += context_item['text'] + '\n\n'
            
            #-------------------------
            # getting expert answer
            #-------------------------
            if use_expert_answer:
                if st.session_state.expert_model in ["LLaMA-TOMMI-1.0-11B", "LLaMA-3.2-11B"]:

                    if st.session_state.expert_model == "LLaMA-TOMMI-1.0-11B":
                        model_ = st.session_state.tommi_model
                        tokenizer_ = st.session_state.tommi_tokenizer
                        
                    elif st.session_state.expert_model == "LLaMA-3.2-11B":
                        model_ = st.session_state.llama_model
                        tokenizer_ = st.session_state.llama_tokenizer
                        
                    messages = [
                        {"role": "system", "content": get_expert_system_prompt()},
                        {"role": "user", "content": st.session_state.question}
                    ]

                    expert_answer = generate_response(
                        model=model_,
                        tokenizer=tokenizer_,
                        messages=messages,
                        tokenizer_max_length=500,
                        do_sample=expert_do_sample,
                        temperature=expert_temperature if expert_do_sample else None,
                        top_k=expert_top_k if expert_do_sample else None,
                        top_p=expert_top_p if expert_do_sample else None,
                        num_beams=expert_num_beams if not expert_do_sample else 1,
                        max_new_tokens=expert_max_new_tokens
                    )

                else:  # openai
                    expert_answer = openai_domain_specific_answer_generation(
                        get_expert_system_prompt(),
                        st.session_state.question,
                        model=st.session_state.expert_model,
                        temperature=expert_api_temperature,
                        top_p=expert_api_top_p
                    )

                st.session_state.expert_answer = fix_latex(expert_answer)

            else:
                st.session_state.expert_answer = 'No Expert Answer. Only use the context.'

            #-------------------------
            # synthesis responses
            #-------------------------
            if st.session_state.synthesis_model in ["LLaMA-3.2-3B", "LLaMA-3.2-11B"]:

                if st.session_state.expert_model == "LLaMA-3.2-11B":
                    model_s = st.session_state.llama_model
                    tokenizer_s = st.session_state.llama_tokenizer
                    
                elif st.session_state.expert_model == "LLaMA-3.2-3B":
                    model_s = st.session_state.llama_model_3B
                    tokenizer_s = st.session_state.llama_tokenizer_3B

                synthesis_prompt = f"""
                Question:
                {st.session_state.question}

                Direct Answer:
                {st.session_state.expert_answer}

                Retrieved Context:
                {context}

                Final Answer:
                """
                messages = [
                    {"role": "system", "content": get_synthesis_system_prompt("Finite Element Method")},
                    {"role": "user", "content": synthesis_prompt}
                ]

                synthesis_answer = generate_response(
                    model=model_s,
                    tokenizer=tokenizer_s,
                    messages=messages,
                    tokenizer_max_length=30000,
                    do_sample=synthesis_do_sample,
                    temperature=synthesis_temperature if synthesis_do_sample else None,
                    top_k=synthesis_top_k if synthesis_do_sample else None,
                    top_p=synthesis_top_p if synthesis_do_sample else None,
                    num_beams=synthesis_num_beams if not synthesis_do_sample else 1,
                    max_new_tokens=synthesis_max_new_tokens
                )

            else:
                synthesis_answer = openai_context_integration(
                    get_synthesis_system_prompt("Finite Element Method"),
                    st.session_state.question,
                    st.session_state.expert_answer,
                    context,
                    model=st.session_state.synthesis_model,
                    temperature=synthesis_api_temperature,
                    top_p=synthesis_api_top_p
                )

        # quick check after getting the answer
        if synthesis_answer.split()[0] == "NOT_ENOUGH_INFO":
            st.markdown("")
            st.markdown("#### Query:")
            st.markdown(fix_latex(st.session_state.question))
            if show_expert_responce:
                st.markdown("#### Initial Expert Answer:")
                st.markdown(st.session_state.expert_answer)
            st.markdown("#### Answer:")
            st.write(":smiling_face_with_tear:")
            st.markdown(synthesis_answer.split('NOT_ENOUGH_INFO')[1])
            st.divider()
            st.caption(get_disclaimer())
            # st.caption("The AI Teaching Assistant project")
            st.session_state.question_answered = False
            st.stop()   
        else:
            st.session_state.answer = fix_latex(synthesis_answer)
            st.session_state.question_answered = True

if st.session_state.question_answered:
    st.markdown("")
    st.markdown("#### Query:")
    st.markdown(fix_latex(st.session_state.question))
    if show_expert_responce:
        st.markdown("#### Initial Expert Answer:")
        st.markdown(st.session_state.expert_answer)
    st.markdown("#### Answer:")
    st.markdown(st.session_state.answer)
    
    if top_k_YT > 0:
        st.markdown("#### Retrieved content in lecture videos")
        for i, (video_id, contexts) in enumerate(st.session_state.context_by_video.items(), start=1):
            # with st.expander(f"**Video {i}** | {contexts[0]['title']}", expanded=True):
            with st.container(border=True):
                st.markdown(f"**Video {i} | {contexts[0]['title']}**")
                video_placeholder = st.empty()
                video_placeholder.markdown(get_youtube_embed(video_id, 0, 0), unsafe_allow_html=True)
                st.markdown('')
                with st.container(border=False):
                    st.markdown("Retrieved Times")
                    cols = st.columns([1 for i in range(len(contexts))] + [9 - len(contexts)])
                    for j, context_item in enumerate(contexts):
                        start_time = int(context_item['start'])
                        label = sec_to_time(start_time)
                        if cols[j].button(label, key=f"{video_id}_{start_time}"):
                            if st.session_state.playing_video_id is not None:
                                st.session_state.playing_video_id = None
                                video_placeholder.empty()
                            video_placeholder.markdown(get_youtube_embed(video_id, start_time, 1), unsafe_allow_html=True)
                            st.session_state.playing_video_id = video_id
                
                with st.expander("Video Summary", expanded=False):
                    # st.write("##### Video Overview:")
                    st.markdown(summary[video_id])

    if show_textbook and top_k_Latex > 0:
        st.markdown("#### Retrieved content in textbook",help="The Finite Element Method: Linear Static and Dynamic Finite Element Analysis")
        for i, (section_id, contexts) in enumerate(st.session_state.context_by_section.items(), start=1):
            # with st.expander(f"**Section {i} | {section_id}**", expanded=True):
            st.markdown(f"**Section {i} | {section_id}**")
            for context_item in contexts:
                st.markdown(context_item['text'])
                st.divider()

st.markdown(" ")
st.divider()
st.caption(get_disclaimer())