import streamlit as st
import os
from datasets import load_dataset
import pandas as pd
import matplotlib.pyplot as plt

HF_TOKEN = os.environ.get("HF_TOKEN")

ds = load_dataset("CohereForAI/mmlu-translations-results", split="train", token=HF_TOKEN)

df = ds.to_pandas()

st.title("MMLU Translations Progress")

# Extract the language from the metadata column and create a new column
df['language'] = df['metadata'].apply(lambda x: x.get('language'))

# Count the occurrences of each language
language_counts = df['language'].value_counts()

# Plotting the bar chart using matplotlib
fig, ax = plt.subplots()
language_counts.plot(kind='bar', ax=ax)
ax.set_title('Number of Completed Tasks for Each Language')
ax.set_xlabel('Language')
ax.set_ylabel('Count')

# Convert the language counts to a DataFrame for display in the table
language_counts_df = language_counts.reset_index()
language_counts_df.columns = ['Language', 'Count']

# Display the table in the Streamlit app
st.table(language_counts_df)

# Display the plot in the Streamlit app
st.pyplot(fig)

st.dataframe(df)