app.py

import os
import streamlit as st
import json
import tarfile

st.set_page_config(layout="wide")

PARENT_DIR: str = os.path.join(os.path.dirname(os.path.abspath(__file__)))
EVAL_DIR: str = os.path.join(PARENT_DIR, "eval-results")
st.title("K2 Evaluation Gallery")
st.markdown("""The K2 gallery allows one to browse the output of various evaluations on intermediate K2 checkpoints, which provides an intuitive understanding on how the model develops and improves over time.""")


with st.sidebar:
    html = f"<img src='https://www.llm360.ai/images/logo-highres.png' width='100' /><img src='https://huggingface.co/spaces/LLM360/k2-eval-gallery/raw/main/k2-logo.svg' width='100' />"
    st.markdown(html, unsafe_allow_html=True)

    metric = st.radio(
        "Choose a metric", options=os.listdir(os.path.join(EVAL_DIR))
    )

    n_shot = st.radio(
        "Selece an n-shot number", os.listdir(os.path.join(EVAL_DIR, metric))
    )

col1, col2 = st.columns(2)

def render_column(col_label):
    st.header(f"Checkpoint {col_label}")
    ckpt = st.select_slider('Select a checkpoint', sorted(os.listdir(os.path.join(EVAL_DIR, metric, n_shot))), key=col_label + '1')
    st.write(f'Veiwing Evaluation Results for Checkpoint: `{ckpt}`')
    file = st.selectbox("Select a file", sorted(f_name[:-len(".tar.gz")] for f_name in os.listdir(os.path.join(EVAL_DIR, metric, n_shot, ckpt))), key=col_label + '2')
    file += ".tar.gz"
    with tarfile.open(os.path.join(EVAL_DIR, metric, n_shot, ckpt, file), "r:gz") as tar:
        f = tar.extractfile(tar.next())
        eval_json = json.load(f)
        if isinstance(eval_json, list):
            doc_id = st.slider("Select a document id", 0, len(eval_json) - 1, 0, 1, key=col_label + '3')
            st.json(eval_json[doc_id])
        else:
            st.json(eval_json)
        f.close()

with col1:
    render_column('A')

with col2:
    render_column('B')