Benjamin Bossan
Initial commit
76ffd6d
raw history blame
No virus
11.1 kB
# HF space creator starting from an sklearn model
import base64
import glob
import io
import json
import os
import pickle
import re
import shutil
from pathlib import Path
from tempfile import mkdtemp
import pandas as pd
import sklearn
import streamlit as st
from sklearn.base import BaseEstimator
import skops.io as sio
from skops import card, hub_utils
st.set_page_config(layout="wide")
st.title("Skops space creator for sklearn")
PLACEHOLDER = "[More Information Needed]"
PLOT_PREFIX = "__plot__:"
custom_sections: dict[str, str] = {}
tmp_repo = Path(mkdtemp(prefix="skops-"))
left_col, right_col = st.columns([1, 2])
# a hacky way to "persist" custom sections
CUSTOM_SECTIONS_CACHE_FILE = ".custom-sections.json"
def _clear_custom_section_cache():
with open(CUSTOM_SECTIONS_CACHE_FILE, "w") as f:
f.write("")
def _load_custom_section_cache():
global custom_sections
# in case file doesn't exist yet, create it
if not os.path.exists(CUSTOM_SECTIONS_CACHE_FILE):
Path(CUSTOM_SECTIONS_CACHE_FILE).touch()
with open(CUSTOM_SECTIONS_CACHE_FILE, "r") as f:
try:
custom_sections = json.load(f)
except ValueError:
pass
def _write_custom_section_cache():
with open(CUSTOM_SECTIONS_CACHE_FILE, "w") as f:
json.dump(custom_sections, f)
def _remove_custom_section(key):
del custom_sections[key]
_write_custom_section_cache()
def _clear_repo(path):
for file_path in glob.glob(str(Path(path) / "*")):
if os.path.isfile(file_path) or os.path.islink(file_path):
os.unlink(file_path)
elif os.path.isdir(file_path):
shutil.rmtree(file_path)
def _write_plot(plot_name, plot_file):
with open(plot_name, "wb") as f:
f.write(plot_file)
def init_repo():
_clear_repo(tmp_repo)
try:
file_name = Path(mkdtemp(prefix="skops-")) / "model.skops"
sio.dump(model, file_name)
hub_utils.init(
model=file_name,
dst=tmp_repo,
task=task,
data=data,
requirements=requirements,
)
except Exception as exc:
print("Uh oh, something went wrong when initializing the repo:", exc)
def load_model():
if model_file is None:
return
bytes_data = model_file.getvalue()
model = pickle.loads(bytes_data)
assert isinstance(model, BaseEstimator), "model must be an sklearn model"
return model
def load_data():
if data_file is None:
return
bytes_data = io.BytesIO(data_file.getvalue())
df = pd.read_csv(bytes_data)
return df
def _parse_metrics(metrics):
metrics_table = {}
for line in metrics.splitlines():
line = line.strip()
name, _, val = line.partition("=")
try:
# try to coerce to float but don't error if it fails
val = float(val.strip())
except ValueError:
pass
metrics_table[name.strip()] = val
return metrics_table
def _create_model_card():
if model is None or data is None:
st.text("*some data is missing to render the model card*")
return
init_repo()
metadata = card.metadata_from_config(tmp_repo)
model_card = card.Card(model=model, metadata=metadata)
if model_description:
model_card.add(**{"Model description": model_description})
if intended_uses:
model_card.add(
**{"Model description/Intended uses & limitations": intended_uses}
)
if metrics:
metrics_table = _parse_metrics(metrics)
model_card.add_metrics(**metrics_table)
if authors:
model_card.add(**{"Model Card Authors": authors})
if contact:
model_card.add(**{"Model Card Contact": contact})
if citation:
model_card.add(**{"Citation": citation})
if custom_sections:
for key, val in custom_sections.items():
if not key:
continue
if key.startswith(PLOT_PREFIX):
key = key[len(PLOT_PREFIX):]
model_card.add_plot(**{key: val})
else:
model_card.add(**{key: val})
return model_card
def _process_card_for_rendering(rendered: str) -> tuple[str, str]:
idx = rendered[1:].index("\n---") + 1
metadata = rendered[3:idx]
rendered = rendered[idx + 4 :] # noqa: E203
# below is a hack to display the images in streamlit
# https://discuss.streamlit.io/t/image-in-markdown/13274/10 The problem is
# that streamlit does not display images in markdown, so we need to replace
# them with html. However, we only want that in the rendered markdown, not
# in the card that is produced for the hub
def markdown_images(markdown):
# example image markdown:
# ![Test image](images/test.png "Alternate text")
images = re.findall(
r'(!\[(?P<image_title>[^\]]+)\]\((?P<image_path>[^\)"\s]+)\s*([^\)]*)\))',
markdown
)
return images
def img_to_bytes(img_path):
img_bytes = Path(img_path).read_bytes()
encoded = base64.b64encode(img_bytes).decode()
return encoded
def img_to_html(img_path, img_alt):
img_format = img_path.split(".")[-1]
img_html = (
f'<img src="data:image/{img_format.lower()};'
f'base64,{img_to_bytes(img_path)}" '
f'alt="{img_alt}" '
'style="max-width: 100%;">'
)
return img_html
def markdown_insert_images(markdown):
images = markdown_images(markdown)
for image in images:
image_markdown = image[0]
image_alt = image[1]
image_path = image[2]
markdown = markdown.replace(image_markdown, img_to_html(image_path, image_alt))
return markdown
rendered_with_img = markdown_insert_images(rendered)
return metadata, rendered_with_img
def display_model_card():
model_card = _create_model_card()
if not model_card:
return
rendered = model_card.render()
metadata, rendered = _process_card_for_rendering(rendered)
# idx = rendered[1:].index("\n---") + 1
# metadata = rendered[3:idx]
# rendered = rendered[idx + 4 :] # noqa: E203
with right_col:
# strip metadata
with st.expander("show metadata"):
st.text(metadata)
st.markdown(rendered, unsafe_allow_html=True)
def download_model_card():
model_card = _create_model_card()
if model_card is not None:
return model_card.render()
return ""
def add_custom_section():
# this is required to "refresh" these variables...
global section_name, section_content
section_name = st.session_state.key_section_name
section_content = st.session_state.key_section_content
if not section_name or not section_content:
return
custom_sections[section_name] = section_content
_write_custom_section_cache()
def add_custom_plot():
# this is required to "refresh" these variables...
global section_name, section_content
plot_name = st.session_state.key_plot_name
plot_file = st.session_state.key_plot_file
if not plot_name or not plot_file:
return
# store plot in temp repo
file_name = plot_file.name.replace(" ", "_")
file_path = str(tmp_repo / file_name)
with open(file_path, "wb") as f:
f.write(plot_file.getvalue())
custom_sections[str(PLOT_PREFIX + plot_name)] = file_path
_write_custom_section_cache()
with left_col:
# This contains every element required to edit the model card
model = None
data = None
section_name = None
section_content = None
model_file = st.file_uploader("Upload a model*", on_change=load_model)
data_file = st.file_uploader(
"Upload X data (csv)*", type=["csv"], on_change=load_data
)
task = st.selectbox(
label="Choose the task type*",
options=[
"tabular-classification",
"tabular-regression",
"text-classification",
"text-regression",
],
on_change=init_repo,
)
requirements = st.text_input(
label="Requirements*",
value=[f"scikit-learn=={sklearn.__version__}\n"],
on_change=init_repo,
)
if model_file is not None:
model = load_model()
if data_file is not None:
data = load_data()
if model is not None and data is not None:
init_repo()
model_description = st.text_input("Model description", value=PLACEHOLDER)
intended_uses = st.text_area(
"Intended uses & limitations", height=2, value=PLACEHOLDER
)
metrics = st.text_area("Metrics (e.g. 'accuracy = 0.95'), one metric per line")
authors = st.text_area(
"Authors",
value="This model card is written by following authors:\n\n" + PLACEHOLDER,
)
contact = st.text_area(
"Contact",
value="You can contact the model card authors through following channels:\n\n"
+ PLACEHOLDER,
)
citation = st.text_area(
"Citation",
value="Below you can find information related to citation.\n\nBibTex:\n\n```\n"
+ PLACEHOLDER
+ "\n```",
height=5,
)
with st.form("custom-section", clear_on_submit=True):
section_name = st.text_input(
"Section name (use '/' for subsections, e.g. 'Model description/My new"
" section')",
key="key_section_name",
)
section_content = st.text_area(
"Content of the new section", key="key_section_content"
)
submit_new_section = st.form_submit_button(
"Create new section", on_click=add_custom_section
)
with st.form("custom-plots", clear_on_submit=True):
plot_name = st.text_input(
"Section name (use '/' for subsections, e.g. 'Model description/My new"
" plot')",
key="key_plot_name",
)
plot_file = st.file_uploader(
"Upload a figure*", key="key_plot_file"
)
submit_new_plot = st.form_submit_button(
"Add plot", on_click=add_custom_plot
)
_load_custom_section_cache()
for key in custom_sections:
if not key:
continue
if key.startswith(PLOT_PREFIX):
st.button(
f"Remove plot '{key[len(PLOT_PREFIX):]}'", on_click=_remove_custom_section, args=(key,)
)
else:
st.button(
f"Remove section '{key}'", on_click=_remove_custom_section, args=(key,)
)
if custom_sections:
st.button(
f"Remove all ({len(custom_sections)}) custom elements",
on_click=_clear_custom_section_cache,
)
with right_col:
# this contains the rendered model card
st.button(label="Render model card", on_click=display_model_card)
rendered = download_model_card()
if rendered:
st.download_button(label="Download model card (markdown format)", data=rendered)