log usage

app.py

@@ -1,10 +1,13 @@
+import csv
 import os
 import shutil
+from datetime import datetime
 from pathlib import Path
 from tempfile import TemporaryDirectory
+from typing import Optional
 
 import gradio as gr
-from huggingface_hub import HfApi, ModelCard, scan_cache_dir
+from huggingface_hub import HfApi, ModelCard, Repository, scan_cache_dir
 from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
 
 from convert import convert
@@ -15,6 +18,22 @@ try:
 except:
     MAX_REPO_SIZE = 24 * 1000 * 1000 * 1000
 
+# Used to log Space usage
+# Taken from https://huggingface.co/spaces/onnx/export
+DATASET_REPO_ID = "Wauplin/bloom.cpp-converters"
+DATASET_LOCAL_DIR = "usage_data"
+DATASET_LOCAL_FILE = Path(DATASET_LOCAL_DIR) / "data" / "data.csv"
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+repo: Optional[Repository] = None
+if HF_TOKEN:
+    repo = Repository(
+        local_dir=DATASET_LOCAL_DIR,
+        clone_from=DATASET_REPO_ID,
+        repo_type="dataset",
+        token=HF_TOKEN,
+    )
+
 
 class Generator:
     # Taken from https://stackoverflow.com/a/34073559
@@ -29,6 +48,14 @@ class Generator:
 def run(
     token: str, model_id: str, precision: str, quantization: bool, destination: str
 ):
+    _log_usage(
+        status="start",
+        model_id=model_id,
+        precision=precision,
+        quantization=quantization,
+        destination=destination,
+        pr_url=None,
+    )
     _all_logs = []
 
     def _log(msg: str):
@@ -38,6 +65,14 @@ def run(
 
     if token == "" or model_id == "":
         yield _log("### Invalid input 🐞\n\nPlease fill a token and model_id.")
+        _log_usage(
+            status="invalid input",
+            model_id=model_id,
+            precision=precision,
+            quantization=quantization,
+            destination=destination,
+            pr_url=None,
+        )
         return
     if destination == "":
         _log("Destination not provided. Will default to the initial repo.")
@@ -52,6 +87,14 @@ def run(
         yield _log(
             f"\n### Error 😢😢😢\n\nRepository {model_id} not found. Only public models are convertible at the moment."
         )
+        _log_usage(
+            status="model not found",
+            model_id=model_id,
+            precision=precision,
+            quantization=quantization,
+            destination=destination,
+            pr_url=None,
+        )
         return
 
     try:
@@ -64,6 +107,14 @@ def run(
             yield _log(
                 f"### Unprocessable 😢😢😢\n\nModel {model_id} is too big and cannot be processed in this Space. This Space needs to be able to load the model in memory before converting it. To avoid a memory issue, we do not process models bigger than {MAX_REPO_SIZE}b.\n\nYou have 2 options:\n- [Duplicate this Space](https://huggingface.co/spaces/Wauplin/bloomz.cpp-converter?duplicate=true) and assign a bigger machine. You will need to set 'MAX_REPO_SIZE' as a secret to overwrite the default value. Once you are done, remove the upgraded hardware and/or delete the Space.\n- Manually convert the weights by following [this guide](https://github.com/NouamaneTazi/bloomz.cpp#usage)."
             )
+            _log_usage(
+                status="unprocessable",
+                model_id=model_id,
+                precision=precision,
+                quantization=quantization,
+                destination=destination,
+                pr_url=None,
+            )
             return
 
         with TemporaryDirectory() as cache_folder:
@@ -130,10 +181,26 @@ def run(
             yield _log(
                 f"### Success 🔥\n\nYay! This model was successfully converted! Make sure to let the repo owner know about it and review your PR. You might need to complete the PR manually, especially to add information in the model card."
             )
+            _log_usage(
+                status="success",
+                model_id=model_id,
+                precision=precision,
+                quantization=quantization,
+                destination=destination,
+                pr_url=pr_url,
+            )
             shutil.rmtree(model_path.parent)
             _delete_cache()
             return
     except Exception as e:
+        _log_usage(
+            status="error",
+            model_id=model_id,
+            precision=precision,
+            quantization=quantization,
+            destination=destination,
+            pr_url=None,
+        )
         yield _log(f"### Error 😢😢😢\n\n{e}")
         _delete_cache()
         return
@@ -147,6 +214,18 @@ def _delete_cache():
     )
 
 
+def _log_usage(**kwargs):
+    # save in a private dataset
+    # Taken from https://huggingface.co/spaces/onnx/export
+    if repo is not None:
+        repo.git_pull(rebase=True)
+        with DATASET_LOCAL_FILE.open("a") as csv_file:
+            writer = csv.DictWriter(csv_file, fieldnames=["time"] + list(kwargs.keys()))
+            writer.writerow({"time": str(datetime.now()), **kwargs})
+        commit_url = repo.push_to_hub()
+        print("[dataset]", commit_url)
+
+
 TITLE = """
 <h1 style="font-weight: 900; font-size: 32px; margin-bottom: 10px; margin-top: 10px; text-align: center;">
     Make any BLOOM-like model compatible with bloomz.cpp