src/expert.py

import streamlit as st
from ecologits.impacts.llm import compute_llm_impacts

from src.utils import format_impacts, average_range_impacts
from src.impacts import display_impacts
from src.electricity_mix import (
    COUNTRY_CODES,
    find_electricity_mix,
    dataframe_electricity_mix,
)
from src.models import load_models
from src.constants import PROMPTS

import plotly.express as px


def reset_model():
    model = "CUSTOM"


def expert_mode():
    st.markdown("### 🤓 Expert mode")

    with st.container(border=True):
        ########## Model info ##########

        col1, col2, col3 = st.columns(3)

        df = load_models(filter_main=True)

        with col1:
            provider_exp = st.selectbox(
                label="Provider",
                options=[x for x in df["provider_clean"].unique()],
                index=7,
                key=1,
            )

        with col2:
            model_exp = st.selectbox(
                label="Model",
                options=[
                    x
                    for x in df["name_clean"].unique()
                    if x
                    in df[df["provider_clean"] == provider_exp]["name_clean"].unique()
                ],
                key=2,
            )

        with col3:
            output_tokens_exp = st.selectbox(
                label="Example prompt", options=[x[0] for x in PROMPTS], key=3
            )

        df_filtered = df[
            (df["provider_clean"] == provider_exp) & (df["name_clean"] == model_exp)
        ]

        try:
            total_params = int(df_filtered["total_parameters"].iloc[0])
        except:
            total_params = int(
                (
                    df_filtered["total_parameters"].values[0]["min"]
                    + df_filtered["total_parameters"].values[0]["max"]
                )
                / 2
            )

        try:
            active_params = int(df_filtered["active_parameters"].iloc[0])
        except:
            active_params = int(
                (
                    df_filtered["active_parameters"].values[0]["min"]
                    + df_filtered["active_parameters"].values[0]["max"]
                )
                / 2
            )

        ########## Model parameters ##########

        col11, col22, col33 = st.columns(3)

        with col11:
            active_params = st.number_input(
                "Active parameters (B)", 0, None, active_params
            )

        with col22:
            total_params = st.number_input(
                "Total parameters (B)", 0, None, total_params
            )

        with col33:
            output_tokens = st.number_input(
                label="Output completion tokens",
                min_value=0,
                value=[x[1] for x in PROMPTS if x[0] == output_tokens_exp][0],
            )

        ########## Electricity mix ##########

        location = st.selectbox("Location", [x[0] for x in COUNTRY_CODES])

        col4, col5, col6 = st.columns(3)

        with col4:
            mix_gwp = st.number_input(
                "Electricity mix - GHG emissions [kgCO2eq / kWh]",
                find_electricity_mix(
                    [x[1] for x in COUNTRY_CODES if x[0] == location][0]
                )[2],
                format="%0.6f",
            )
            # disp_ranges = st.toggle('Display impact ranges', False)
        with col5:
            mix_adpe = st.number_input(
                "Electricity mix - Abiotic resources [kgSbeq / kWh]",
                find_electricity_mix(
                    [x[1] for x in COUNTRY_CODES if x[0] == location][0]
                )[0],
                format="%0.13f",
            )
        with col6:
            mix_pe = st.number_input(
                "Electricity mix - Primary energy [MJ / kWh]",
                find_electricity_mix(
                    [x[1] for x in COUNTRY_CODES if x[0] == location][0]
                )[1],
                format="%0.3f",
            )

    impacts = compute_llm_impacts(
        model_active_parameter_count=active_params,
        model_total_parameter_count=total_params,
        output_token_count=output_tokens,
        request_latency=100000,
        if_electricity_mix_gwp=mix_gwp,
        if_electricity_mix_adpe=mix_adpe,
        if_electricity_mix_pe=mix_pe,
    )

    impacts, usage, embodied = format_impacts(impacts)

    with st.container(border=True):
        st.markdown(
            '<h3 align="center">Environmental Impacts</h2>', unsafe_allow_html=True
        )

        display_impacts(impacts)

    with st.expander("⚖️ Usage vs Embodied"):
        st.markdown(
            '<h3 align="center">Embodied vs Usage comparison</h2>',
            unsafe_allow_html=True,
        )

        st.markdown(
            "The usage impacts account for the electricity consumption of the model while the embodied impacts account for resource extraction (e.g., minerals and metals), manufacturing, and transportation of the hardware."
        )

        col_ghg_comparison, col_adpe_comparison, col_pe_comparison = st.columns(3)

        with col_ghg_comparison:
            fig_gwp = px.pie(
                values=[
                    average_range_impacts(usage.gwp.value),
                    average_range_impacts(embodied.gwp.value),
                ],
                names=["usage", "embodied"],
                title="GHG emissions",
                color_discrete_sequence=["#00BF63", "#0B3B36"],
                width=100,
            )
            fig_gwp.update_layout(showlegend=False, title_x=0.5)

            st.plotly_chart(fig_gwp)

        with col_adpe_comparison:
            fig_adpe = px.pie(
                values=[
                    average_range_impacts(usage.adpe.value),
                    average_range_impacts(embodied.adpe.value),
                ],
                names=["usage", "embodied"],
                title="Abiotic depletion",
                color_discrete_sequence=["#0B3B36", "#00BF63"],
                width=100,
            )
            fig_adpe.update_layout(showlegend=False, title_x=0.5)

            st.plotly_chart(fig_adpe)

        with col_pe_comparison:
            fig_pe = px.pie(
                values=[
                    average_range_impacts(usage.pe.value),
                    average_range_impacts(embodied.pe.value),
                ],
                names=["usage", "embodied"],
                title="Primary energy",
                color_discrete_sequence=["#00BF63", "#0B3B36"],
                width=100,
            )
            fig_pe.update_layout(showlegend=False, title_x=0.5)

            st.plotly_chart(fig_pe)

    with st.expander("🌍️ Location impact"):
        st.markdown(
            '<h4 align="center">How can location impact the footprint ?</h4>',
            unsafe_allow_html=True,
        )

        countries_to_compare = st.multiselect(
            label="Countries to compare",
            options=[x[0] for x in COUNTRY_CODES],
            default=["🇫🇷 France", "🇺🇸 United States", "🇨🇳 China"],
        )

        try:
            df_comp = dataframe_electricity_mix(countries_to_compare)

            impact_type = st.selectbox(
                label="Select an impact type to compare",
                options=[x for x in df_comp.columns if x != "country"],
                index=1,
            )

            df_comp.sort_values(by=impact_type, inplace=True)

            fig_2 = px.bar(
                df_comp,
                x=df_comp.index,
                y=impact_type,
                text=impact_type,
                color=impact_type,
            )

            st.plotly_chart(fig_2)

        except:
            st.warning("Can't display chart with no values.")