Add cache.

app.py

@@ -6,27 +6,36 @@ from core.constants import OAUTH_CLIENT_ID
 from core.constants import OAUTH_STATE
 from core.constants import REDIRECT_URI
 from core.state import CurrentStep
+from core.state import get_cached_user
 from core.state import User
 from utils import init_state
 from views.splash import render_splash
 from views.wizard import render_editor
 
-init_state()
-
-
 st.set_page_config(page_title="Croissant Editor", page_icon="🥐", layout="wide")
-col1, col2 = st.columns([10, 1])
+col1, col2, col3 = st.columns([10, 1, 1])
 col1.header("Croissant Editor")
 
-if OAUTH_CLIENT_ID and not st.session_state.get(User):
+init_state()
+
+user = get_cached_user()
+
+if OAUTH_CLIENT_ID and not user:
     query_params = st.experimental_get_query_params()
     state = query_params.get("state")
     if state and state[0] == OAUTH_STATE:
         code = query_params.get("code")
         if not code:
             st.stop()
-        st.session_state[User] = User.connect(code)
-        st.experimental_set_query_params()
+        try:
+            st.session_state[User] = User.connect(code)
+            # Clear the cache to force retrieving the new user.
+            get_cached_user.clear()
+            get_cached_user()
+        except:
+            raise
+        finally:
+            st.experimental_set_query_params()
     else:
         redirect_uri = urllib.parse.quote(REDIRECT_URI, safe="")
         client_id = urllib.parse.quote(OAUTH_CLIENT_ID, safe="")
@@ -42,10 +51,18 @@ def _back_to_menu():
     init_state(force=True)
 
 
+def _logout():
+    """Logs the user out."""
+    st.cache_data.clear()
+
+
 if st.session_state[CurrentStep] != CurrentStep.splash:
     col2.write("\n")  # Vertical box to shift the button menu
     col2.button("Menu", on_click=_back_to_menu)
 
+col3.write("\n")  # Vertical box to shift the lgout menu
+col3.button("Log out", on_click=_logout)
+
 
 if st.session_state[CurrentStep] == CurrentStep.splash:
     render_splash()

core/past_projects.py

@@ -6,12 +6,12 @@ import streamlit as st
 
 from core.constants import PAST_PROJECTS_PATH
 from core.state import CurrentProject
+from core.state import get_cached_user
 from core.state import Metadata
-from core.state import User
 
 
 def load_past_projects_paths() -> list[epath.Path]:
-    user = st.session_state.get(User)
+    user = get_cached_user()
     past_projects_path = PAST_PROJECTS_PATH(user)
     past_projects_path.mkdir(parents=True, exist_ok=True)
     return sorted(list(past_projects_path.iterdir()), reverse=True)

core/state.py

@@ -8,7 +8,6 @@ from __future__ import annotations
 import base64
 import dataclasses
 import datetime
-import hashlib
 from typing import Any
 
 from etils import epath
@@ -64,17 +63,30 @@ class User:
             access_token = response.get("access_token")
             id_token = response.get("id_token")
             if access_token and id_token:
-                # Warning: this is temporary while being able to retrieve the username.
-                username = hashlib.sha256(access_token.encode()).hexdigest()
-                return User(
-                    access_token=access_token, username=username, id_token=id_token
-                )
+                url = "https://huggingface.co/oauth/userinfo"
+                headers = {"Authorization": f"Bearer {access_token}"}
+                response = requests.get(url, headers=headers)
+                if response.status_code == 200:
+                    response = response.json()
+                    username = response.get("preferred_username")
+                    if username:
+                        return User(
+                            access_token=access_token,
+                            username=username,
+                            id_token=id_token,
+                        )
         raise Exception(
             f"Could not connect to Hugging Face. Please, go to {REDIRECT_URI}."
             f" ({response=})."
         )
 
 
+@st.cache_data(ttl=datetime.timedelta(hours=1))
+def get_cached_user():
+    """Caches user in session_state."""
+    return st.session_state.get(User)
+
+
 class CurrentStep:
     """Holds all major state variables for the application."""
 
@@ -91,8 +103,8 @@ class CurrentProject:
     @classmethod
     def create_new(cls) -> CurrentProject:
         timestamp = datetime.datetime.now().strftime(PROJECT_FOLDER_PATTERN)
-        user = st.session_state.get(User)
-        if user is None:
+        user = get_cached_user()
+        if user is None and OAUTH_CLIENT_ID:
             return None
         else:
             path = PAST_PROJECTS_PATH(user)

views/previous_files.py

@@ -29,12 +29,12 @@ def render_previous_files():
     else:
         for index, path in enumerate(paths):
             try:
-                col1, col2 = st.columns([10, 1])
                 metadata = open_project(path)
                 timestamp = datetime.datetime.strptime(
                     path.name, PROJECT_FOLDER_PATTERN
                 ).strftime("%Y/%m/%d %H:%M")
                 label = f"{metadata.name or 'Unnamed dataset'} - {timestamp}"
+                col1, col2 = st.columns([10, 1])
                 col1.button(
                     label,
                     key=f"splash-{index}-load",

views/splash.py

@@ -1,9 +1,13 @@
+import logging
+
+import requests
 import streamlit as st
 
 from core.constants import OAUTH_CLIENT_ID
 from core.state import CurrentProject
 from core.state import CurrentStep
 from core.state import Metadata
+import mlcroissant as mlc
 from utils import jump_to
 from views.load import render_load
 from views.previous_files import render_previous_files
@@ -33,6 +37,29 @@ def render_splash():
                 on_click=create_new_croissant,
                 type="primary",
             )
+        with st.expander("**Try out an example!**", expanded=True):
+
+            def create_example():
+                url = "https://raw.githubusercontent.com/mlcommons/croissant/main/datasets/titanic/metadata.json"
+                try:
+                    json = requests.get(url).json()
+                    metadata = mlc.Metadata.from_json(mlc.Issues(), json, None)
+                    st.session_state[Metadata] = Metadata.from_canonical(metadata)
+                    st.session_state[CurrentProject] = CurrentProject.create_new()
+                    jump_to(CurrentStep.editor)
+                except Exception as exception:
+                    logging.error(exception)
+                    st.write(
+                        "Sorry, it seems that the example is broken... Can you please"
+                        " [open an issue on"
+                        " GitHub](https://github.com/mlcommons/croissant/issues/new)?"
+                    )
+
+            st.button(
+                "Titanic dataset",
+                on_click=create_example,
+                type="primary",
+            )
     with col2:
         with st.expander("**Past projects**", expanded=True):
             render_previous_files()