Add HF repo creation feature

app.py

@@ -4,14 +4,40 @@ This ties together the different parts of the app.
 
 """
 
+from pathlib import Path
+from tempfile import mkdtemp
+from typing import Literal
+
 import streamlit as st
 
-from start import start_input_form
+from create import create_repo_input_form
 from edit import edit_input_form
+from start import start_input_form
+
+
+# Create a hf_path, which is where the repo will be created locally. When the
+# session is created, copy the dummy cat.png file there and make it the cwd
+if "hf_path" not in st.session_state:
+    hf_path = Path(mkdtemp(prefix="skops-"))
+    st.session_state.hf_path = hf_path
+
 
 st.header("Skops model card creator")
 
-if not st.session_state.get("model_card"):
+
+class Screen:
+    state: Literal["start", "edit", "create_repo"] = "start"
+
+
+if "screen" not in st.session_state:
+    st.session_state.screen: Screen = Screen()
+
+
+if st.session_state.screen.state == "start":
     start_input_form()
-else:
+elif st.session_state.screen.state == "edit":
     edit_input_form()
+elif st.session_state.screen.state == "create_repo":
+    create_repo_input_form()
+else:
+    st.write("Something went wrong, please open an issue")

create.py

@@ -0,0 +1,129 @@
+import os
+from pathlib import Path
+
+import streamlit as st
+from skops import hub_utils
+
+from utils import get_rendered_model_card
+
+
+def _add_back_button():
+    def fn():
+        st.session_state.screen.state = "edit"
+
+    st.button("Back", help="continue editing the model card", on_click=fn)
+
+
+def _add_delete_button():
+    def fn():
+        st.session_state.screen.state = "start"
+        if "model_card" in st.session_state:
+            del st.session_state["model_card"]
+        if "task_state" in st.session_state:
+            st.session_state.task_state.reset()
+        if "create_repo_name" in st.session_state:
+            del st.session_state["create_repo_name"]
+        if "hf_token" in st.session_state:
+            del st.session_state["hf_token"]
+
+    st.button("Delete", on_click=fn, help="Start over from scratch (lose all progress)")
+
+
+def _save_model_card(path: Path) -> None:
+    model_card = st.session_state.get("model_card")
+    if model_card:
+        # do not use model_card.save, see doc of get_rendered_model_card
+        rendered = get_rendered_model_card(
+            model_card, hf_path=str(st.session_state.hf_path)
+        )
+        with open(path / "README.md", "w") as f:
+            f.write(rendered)
+
+
+def _display_repo_overview(path: Path) -> None:
+    text = "Files included in the repository:\n"
+    for file in os.listdir(path):
+        size = os.path.getsize(path / file)
+        text += f"- `{file} ({size:,} bytes)`\n"
+    st.markdown(text)
+
+
+def _display_private_box():
+    tip = (
+        "Private repositories can only seen by you or members of the same "
+        "organization, see https://huggingface.co/docs/hub/repositories-settings"
+    )
+    st.checkbox("Make repo private", value=True, help=tip, key="create_repo_private")
+
+
+def _repo_id_field():
+    st.text_input("Name of the repository (e.g. 'User/MyRepo')", key="create_repo_name")
+
+
+def _hf_token_field():
+    tip = "The Hugging Face token can be found at https://hf.co/settings/token"
+    st.text_input(
+        "Enter your Hugging Face token ('hf_***')", key="hf_token", help=tip
+    )
+
+
+def _create_hf_repo(path, repo_name, hf_token, private):
+    try:
+        hub_utils.push(
+            repo_id=repo_name,
+            source=path,
+            token=hf_token,
+            private=private,
+            create_remote=True,
+        )
+    except Exception as exc:
+        st.error(
+            "Oops, something went wrong, please create an issue. "
+            f"The error message is:\n\n{exc}"
+        )
+        return
+
+    st.success(f"Successfully created the repo 'https://huggingface.co/{repo_name}'")
+
+
+def _add_create_repo_button():
+    private = bool(st.session_state.get("create_repo_private"))
+    repo_name = st.session_state.get("create_repo_name")
+    hf_token = st.session_state.get("hf_token")
+    disabled = (not repo_name) or (not hf_token)
+
+    button_text = "Create a new repository"
+    tip = "Creating a repo requires a name and a token"
+    path = st.session_state.get("hf_path")
+    st.button(
+        button_text,
+        help=tip,
+        disabled=disabled,
+        on_click=_create_hf_repo,
+        args=(path, repo_name, hf_token, private),
+    )
+
+    if not repo_name:
+        st.info("Repository name is required")
+    if not hf_token:
+        st.info("Token is required")
+
+
+def create_repo_input_form():
+    if not st.session_state.screen.state == "create_repo":
+        return
+
+    col_0, col_1, *_ = st.columns([2, 2, 2, 2])
+    with col_0:
+        _add_back_button()
+    with col_1:
+        _add_delete_button()
+
+    hf_path = st.session_state.hf_path
+    _save_model_card(hf_path)
+    _display_repo_overview(hf_path)
+    _display_private_box()
+    st.markdown("---")
+    _repo_id_field()
+    _hf_token_field()
+    _add_create_repo_button()

edit.py

@@ -35,7 +35,11 @@ from huggingface_hub import hf_hub_download
 from skops import card
 from skops.card._model_card import PlotSection, split_subsection_names
 
-from utils import iterate_key_section_content, process_card_for_rendering
+from utils import (
+    get_rendered_model_card,
+    iterate_key_section_content,
+    process_card_for_rendering,
+)
 from tasks import (
     AddMetricsTask,
     AddSectionTask,
@@ -94,17 +98,20 @@ def _update_model_card(
         return
 
     if is_fig:
-        fpath = None
+        old_path, fpath = None, None
         if new_content:  # new figure uploaded
             fname = new_content.name.replace(" ", "_")
-            fpath = tmp_path / fname
+            fpath = st.session_state.hf_path / fname
+            old_path = fpath.parent / model_card.select(key).content.path
+
         task = UpdateFigureTask(
             model_card,
             key=key,
             old_name=section_name,
             new_name=new_title,
             data=new_content,
-            path=fpath,
+            new_path=fpath,
+            old_path=old_path,
         )
     else:
         task = UpdateSectionTask(
@@ -128,12 +135,17 @@ def _add_section(model_card: card.Card, key: str) -> None:
 
 def _add_figure(model_card: card.Card, key: str) -> None:
     section_name = f"{key}/Untitled"
-    task = AddFigureTask(model_card, title=section_name, content="cat.png")
+    hf_path = st.session_state.hf_path
+    task = AddFigureTask(
+        model_card, path=hf_path, title=section_name, content="cat.png"
+    )
     st.session_state.task_state.add(task)
 
 
-def _delete_section(model_card: card.Card, key: str) -> None:
-    task = DeleteSectionTask(model_card, key=key)
+def _delete_section(
+    model_card: card.Card, key: str, path: Path
+) -> None:
+    task = DeleteSectionTask(model_card, key=key, path=path)
     st.session_state.task_state.add(task)
 
 
@@ -197,10 +209,11 @@ def create_form_from_section(
 
     col_0, col_1, col_2 = st.columns([4, 2, 2])
     with col_0:
+        path = st.session_state.hf_path / content.path if is_fig else None
         st.button(
             f"delete '{arepr.repr(old_title)}'",
             on_click=_delete_section,
-            args=(model_card, key),
+            args=(model_card, key, path),
             key=f"{key}.delete",
             help="Delete this section, including all its subsections",
         )
@@ -267,6 +280,7 @@ def reset_model_card() -> None:
 
 
 def delete_model_card() -> None:
+    st.session_state.screen.state = "start"
     if "model_card" in st.session_state:
         del st.session_state["model_card"]
     if "task_state" in st.session_state:
@@ -284,19 +298,32 @@ def redo_last():
 
 
 def add_download_model_card_button():
-    model_card = st.session_state.get("model_card")
-    download_disabled = not bool(model_card)
-    data = model_card.render()
+    model_card = st.session_state.model_card
+    data = get_rendered_model_card(
+        model_card, hf_path=str(st.session_state.hf_path)
+    )
     tip = "Download the generated model card as markdown file"
     st.download_button(
         "Save (md)",
         data=data,
-        disabled=download_disabled,
         help=tip,
         file_name="README.md",
     )
 
 
+def add_create_repo_button():
+    def fn():
+        st.session_state.screen.state = "create_repo"
+
+    button_disabled = not bool(st.session_state.get("model_card"))
+    st.button(
+        "Create Repo",
+        help="Create a model repository on Hugging Face Hub",
+        on_click=fn,
+        disabled=button_disabled,
+    )
+
+
 def display_edit_buttons():
     # first row: undo + redo + reset
     col_0, col_1, col_2, *_ = st.columns([2, 2, 2, 2])
@@ -314,11 +341,13 @@ def display_edit_buttons():
         tip = "Undo all edits"
         st.button("Reset", on_click=reset_model_card, help=tip)
 
-    # second row: download + delete
-    col_0, col_1, *_ = st.columns([2, 2, 2, 2])
+    # second row: download + create repo + delete
+    col_0, col_1, col_2, *_ = st.columns([2, 2, 2, 2])
     with col_0:
         add_download_model_card_button()
     with col_1:
+        add_create_repo_button()
+    with col_2:
         tip = "Start over from scratch (lose all progress)"
         st.button("Delete", on_click=delete_model_card, help=tip)
 

start.py

@@ -30,7 +30,6 @@ import skops.io as sio
 from skops import card, hub_utils
 
 
-hf_path = Path(mkdtemp(prefix="skops-"))  # hf repo
 tmp_path = Path(mkdtemp(prefix="skops-"))  # temporary files
 description = """Create an sklearn model card
 
@@ -70,7 +69,8 @@ def _clear_repo(path: str) -> None:
             shutil.rmtree(file_path)
 
 
-def init_repo(path: str) -> None:
+def init_repo() -> None:
+    path = st.session_state.hf_path
     _clear_repo(path)
     requirements = []
     task = "tabular-classification"
@@ -104,26 +104,28 @@ def init_repo(path: str) -> None:
 
 
 def create_skops_model_card() -> None:
-    init_repo(hf_path)
-    metadata = card.metadata_from_config(hf_path)
+    init_repo()
+    metadata = card.metadata_from_config(st.session_state.hf_path)
     model_card = card.Card(model=st.session_state.model, metadata=metadata)
     st.session_state.model_card = model_card
     st.session_state.model_card_type = "skops"
+    st.session_state.screen.state = "edit"
 
 
 def create_empty_model_card() -> None:
-    init_repo(hf_path)
-    metadata = card.metadata_from_config(hf_path)
+    init_repo()
+    metadata = card.metadata_from_config(st.session_state.hf_path)
     model_card = card.Card(
         model=st.session_state.model, metadata=metadata, template=None
     )
     model_card.add(**{"Untitled": "[More Information Needed]"})
     st.session_state.model_card = model_card
     st.session_state.model_card_type = "empty"
+    st.session_state.screen.state = "edit"
 
 
 def create_hf_model_card() -> None:
-    repo_id = st.session_state.get("hf_repo_id", "").strip("'").strip('"')
+    repo_id = st.session_state.get("hf_repo_id", "").strip().strip("'").strip('"')
     if not repo_id:
         return
 
@@ -139,6 +141,7 @@ def create_hf_model_card() -> None:
     model_card = card.parse_modelcard(path)
     st.session_state.model_card = model_card
     st.session_state.model_card_type = "loaded"
+    st.session_state.screen.state = "edit"
 
 
 def start_input_form():
@@ -158,7 +161,10 @@ def start_input_form():
         "Upload an sklearn model (strongly recommended)\n"
         "The model can be used to automatically populate fields in the model card."
     )
-    st.file_uploader("Upload a model*", on_change=load_model, key="model_file")
+
+    if not st.session_state.get("model_file"):
+        st.file_uploader("Upload a model*", on_change=load_model, key="model_file")
+
     st.markdown("---")
 
     st.text(
@@ -180,7 +186,6 @@ def start_input_form():
         ],
         key="task",
         on_change=init_repo,
-        args=(hf_path,),
     )
     st.markdown("---")
 
@@ -189,7 +194,6 @@ def start_input_form():
         value=f"scikit-learn=={sklearn.__version__}\n",
         key="requirements",
         on_change=init_repo,
-        args=(hf_path,),
     )
     st.markdown("---")
 

tasks.py

@@ -5,7 +5,9 @@ Tasks are used to implement "undo" and "redo" functionality.
 """
 from __future__ import annotations
 
+import shutil
 from pathlib import Path
+from tempfile import mkdtemp
 from uuid import uuid4
 
 from skops import card
@@ -80,26 +82,42 @@ class AddSectionTask(Task):
 
 
 class AddFigureTask(Task):
-    """Add a new figure section"""
+    """Add a new figure section
+
+    Figure always starts out with dummy image cat.png.
+
+    """
 
     def __init__(
         self,
         model_card: card.Card,
+        path: Path,
         title: str,
         content: str,
     ) -> None:
         self.model_card = model_card
         self.title = title
-        self.key = title + " " + str(uuid4())[:6]
-        self.content = content
+
+        # Create a unique file name, since the same image can exist more than
+        # once per model card.
+        fname = Path(content)
+        stem = fname.stem
+        suffix = fname.suffix
+        uniq = str(uuid4())[:6]
+        new_fname = str(path / stem) + "_" + uniq + suffix
+
+        self.key = title + " " + uniq
+        self.content = Path(new_fname)
 
     def do(self) -> None:
+        shutil.copy("cat.png", self.content)
         self.model_card.add_plot(**{self.key: self.content})
         section = self.model_card.select(self.key)
         section.title = split_subsection_names(self.title)[-1]
         section.is_fig = True  # type: ignore
 
     def undo(self) -> None:
+        self.content.unlink(missing_ok=True)
         self.model_card.delete(self.key)
 
 
@@ -115,15 +133,23 @@ class DeleteSectionTask(Task):
         self,
         model_card: card.Card,
         key: str,
+        path: Path | None,
     ) -> None:
         self.model_card = model_card
         self.key = key
+        # when 'deleting' a file, move it to a temp file
+        self.path = path
+        self.tmp_path = Path(mkdtemp(prefix="skops-")) / str(uuid4())
 
     def do(self) -> None:
         self.model_card.select(self.key).visible = False
+        if self.path:
+            shutil.move(self.path, self.tmp_path)
 
     def undo(self) -> None:
         self.model_card.select(self.key).visible = True
+        if self.path:
+            shutil.move(self.tmp_path, self.path)
 
 
 class UpdateSectionTask(Task):
@@ -159,7 +185,20 @@ class UpdateSectionTask(Task):
 
 
 class UpdateFigureTask(Task):
-    """Change the title or image of a figure section"""
+    """Change the title or image of a figure section
+
+    Changing the title is easy, just replace it and be done.
+
+    Changing the figure is a bit more tricky. The old figure is in the hf_path
+    under its old name. The new figure is an UploadFile object. For the DO
+    operation, move the old figure to a temporary file and store the UploadFile
+    content to a new file (which may have a different name).
+
+    For the UNDO operation, delete the new figure (its content is still stored
+    in the UploadFile) and move back the old figure from its temporary file to
+    the original location (with its original name).
+
+    """
 
     def __init__(
         self,
@@ -168,14 +207,18 @@ class UpdateFigureTask(Task):
         old_name: str,
         new_name: str,
         data: UploadedFile | None,
-        path: Path | None,
+        new_path: Path | None,
+        old_path: Path | None,
     ) -> None:
         self.model_card = model_card
         self.key = key
         self.old_name = old_name
         self.new_name = new_name
         self.old_data = self.model_card.select(self.key).content
-        self.path = path
+        self.new_path = new_path
+        self.old_path = old_path
+        # when 'deleting' the old image, move to temp path
+        self.tmp_path = Path(mkdtemp(prefix="skops-")) / str(uuid4())
 
         if not data:
             self.new_data = self.old_data
@@ -192,12 +235,14 @@ class UpdateFigureTask(Task):
         # write figure
         # note: this can still be the same image if the image is a file, there
         # is no test to check, e.g., the hash of the image
-        with open(self.path, "wb") as f:
+        shutil.move(self.old_path, self.tmp_path)
+
+        with open(self.new_path, "wb") as f:
             f.write(self.new_data.getvalue())
         section.content = PlotSection(
             alt_text=self.new_data.name,
-            path=self.path,
-        ).format()
+            path=self.new_path,
+        )
 
     def undo(self) -> None:
         section = self.model_card.select(self.key)
@@ -206,7 +251,8 @@ class UpdateFigureTask(Task):
         if self.new_data == self.old_data:  # image is same
             return
 
-        self.path.unlink(missing_ok=True)
+        self.new_path.unlink(missing_ok=True)
+        shutil.move(self.tmp_path, self.old_path)
         section.content = self.old_data
 
 

utils.py

@@ -3,14 +3,31 @@
 from __future__ import annotations
 
 import base64
+import os
 import re
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Iterator
 
+from skops import card
 from skops.card._model_card import Section
 
 
+def get_rendered_model_card(model_card: card.Card, hf_path: str) -> str:
+    # This is a bit hacky:
+    # As a space, the model card is created in a temporary hf_path directory,
+    # which is where all the files are put. So e.g. if a figure is added, it is
+    # found at /tmp/skops-jtyqdgk3/fig.png. However, when the model card is is
+    # actually used, we don't want that, since there, the files will be in the
+    # cwd. Therefore, we remove the tmp directory everywhere we find it in the
+    # file.
+    if not hf_path.endswith(os.path.sep):
+        hf_path += os.path.sep
+
+    rendered = model_card.render()
+    rendered = rendered.replace(hf_path, "")
+    return rendered
+
+
 def process_card_for_rendering(rendered: str) -> tuple[str, str]:
     idx = rendered[1:].index("\n---") + 1
     metadata = rendered[3:idx]