big refactor to include original model card

app.py

@@ -3,7 +3,6 @@ import subprocess
 import signal
 import tempfile
 from pathlib import Path
-from textwrap import dedent
 import logging
 import gradio as gr
 from huggingface_hub import HfApi, ModelCard, whoami
@@ -32,12 +31,7 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
-def get_llama_cpp_notes(
-    gguf_files,
-    new_repo_url,
-    split_model,
-    model_id=None,
-):
+def get_llama_cpp_version():
     try:
         result = subprocess.run(
             ["git", "-C", "./llama.cpp", "describe", "--tags", "--always"],
@@ -47,58 +41,13 @@ def get_llama_cpp_notes(
             text=True,
         )
         version = result.stdout.strip().split("-")[0]
-        text = f"""
-*Produced by [Antigma Labs](https://antigma.ai), [Antigma Quantize Space](https://huggingface.co/spaces/Antigma/quantize-my-repo)*
-
-*Follow Antigma Labs in X [https://x.com/antigma_labs](https://x.com/antigma_labs)*
-
-*Antigma's GitHub Homepage [https://github.com/AntigmaLabs](https://github.com/AntigmaLabs)*
-
-## llama.cpp quantization
-Using <a href="https://github.com/ggml-org/llama.cpp">llama.cpp</a> release <a href="https://github.com/ggml-org/llama.cpp/releases/tag/{version}">{version}</a> for quantization.
-Original model: https://huggingface.co/{model_id}
-Run them directly with [llama.cpp](https://github.com/ggml-org/llama.cpp), or any other llama.cpp based project
-## Prompt format
-```
-<｜begin▁of▁sentence｜>{{system_prompt}}<｜User｜>{{prompt}}<｜Assistant｜><｜end▁of▁sentence｜><｜Assistant｜>
-```
-## Download a file (not the whole branch) from below:
-| Filename | Quant type | File Size | Split |
-| -------- | ---------- | --------- | ----- |
-| {'|'.join(['|'.join([gguf_files[i][0][:-5] if split_model else ('['+gguf_files[i][0]+']'+'(' + new_repo_url+'/blob/main/'+gguf_files[i][0] + ')'), gguf_files[i][3], f"{gguf_files[i][2]:.2f}" + ' GB', str(split_model),'''
-''']) for i in range(len(gguf_files))]) } 
-## Downloading using huggingface-cli
-<details>
-  <summary>Click to view download instructions</summary>
-First, make sure you have hugginface-cli installed:
-
-```
-pip install -U "huggingface_hub[cli]"
-```
-
-Then, you can target the specific file you want:
-
-```
-huggingface-cli download {new_repo_url} --include "{gguf_files[0][0]}" --local-dir ./
-```
-
-If the model is bigger than 50GB, it will have been split into multiple files. In order to download them all to a local folder, run:
-
-```
-huggingface-cli download {new_repo_url} --include "{gguf_files[0][0]}/*" --local-dir ./
-```
-
-You can either specify a new local-dir (deepseek-ai_DeepSeek-V3-0324-Q8_0) or download them all in place (./)
-
-</details>
-"""
-        return text
+        return version
     except subprocess.CalledProcessError as e:
-        print("Error:", e.stderr.strip())
+        logger.error("Error getting llama.cpp version: %s", e.stderr.strip())
         return None
 
 
-def get_repo_namespace(repo_owner, username, user_orgs):
+def get_repo_namespace(repo_owner: str, username: str, user_orgs: list) -> str:
     if repo_owner == "self":
         return username
     for org in user_orgs:
@@ -117,7 +66,7 @@ def escape(s: str) -> str:
     )
 
 
-def toggle_repo_owner(export_to_org, oauth_token: gr.OAuthToken | None):
+def toggle_repo_owner(export_to_org: bool, oauth_token: gr.OAuthToken | None) -> tuple:
     if oauth_token is None or oauth_token.token is None:
         raise gr.Error("You must be logged in to use quantize-my-repo")
     if not export_to_org:
@@ -131,7 +80,9 @@ def toggle_repo_owner(export_to_org, oauth_token: gr.OAuthToken | None):
     )
 
 
-def generate_importance_matrix(model_path: str, train_data_path: str, output_path: str):
+def generate_importance_matrix(
+    model_path: str, train_data_path: str, output_path: str
+) -> None:
     imatrix_command = [
         "./llama.cpp/llama-imatrix",
         "-m",
@@ -147,25 +98,27 @@ def generate_importance_matrix(model_path: str, train_data_path: str, output_pat
     ]
 
     if not os.path.isfile(model_path):
-        raise Exception(f"Model file not found: {model_path}")
+        raise FileNotFoundError(f"Model file not found: {model_path}")
 
-    print("Running imatrix command...")
+    logger.info("Running imatrix command...")
     process = subprocess.Popen(imatrix_command, shell=False)
 
     try:
-        process.wait(timeout=60)  # added wait
+        process.wait(timeout=60)
     except subprocess.TimeoutExpired:
-        print(
+        logger.warning(
             "Imatrix computation timed out. Sending SIGINT to allow graceful termination..."
         )
         process.send_signal(signal.SIGINT)
         try:
-            process.wait(timeout=5)  # grace period
+            process.wait(timeout=5)
         except subprocess.TimeoutExpired:
-            print("Imatrix proc still didn't term. Forecfully terming process...")
+            logger.error(
+                "Imatrix proc still didn't term. Forecfully terming process..."
+            )
             process.kill()
 
-    print("Importance matrix generation completed.")
+    logger.info("Importance matrix generation completed.")
 
 
 def split_upload_model(
@@ -173,101 +126,160 @@ def split_upload_model(
     outdir: str,
     repo_id: str,
     oauth_token: gr.OAuthToken | None,
-    split_max_tensors=256,
-    split_max_size=None,
-    org_token=None,
-    export_to_org=False,
-):
-    print(f"Model path: {model_path}")
-    print(f"Output dir: {outdir}")
+    split_max_tensors: int = 256,
+    split_max_size: str | None = None,
+    org_token: str | None = None,
+    export_to_org: bool = False,
+) -> None:
+    logger.info("Model path: %s", model_path)
+    logger.info("Output dir: %s", outdir)
 
     if oauth_token is None or oauth_token.token is None:
         raise ValueError("You have to be logged in.")
 
-    split_cmd = [
-        "./llama.cpp/llama-gguf-split",
-        "--split",
-    ]
+    split_cmd = ["./llama.cpp/llama-gguf-split", "--split"]
     if split_max_size:
-        split_cmd.append("--split-max-size")
-        split_cmd.append(split_max_size)
+        split_cmd.extend(["--split-max-size", split_max_size])
     else:
-        split_cmd.append("--split-max-tensors")
-        split_cmd.append(str(split_max_tensors))
+        split_cmd.extend(["--split-max-tensors", str(split_max_tensors)])
 
-    # args for output
-    model_path_prefix = ".".join(
-        model_path.split(".")[:-1]
-    )  # remove the file extension
-    split_cmd.append(model_path)
-    split_cmd.append(model_path_prefix)
+    model_path_prefix = ".".join(model_path.split(".")[:-1])
+    split_cmd.extend([model_path, model_path_prefix])
 
-    print(f"Split command: {split_cmd}")
+    logger.info("Split command: %s", split_cmd)
 
     result = subprocess.run(split_cmd, shell=False, capture_output=True, text=True)
-    print(f"Split command stdout: {result.stdout}")
-    print(f"Split command stderr: {result.stderr}")
+    logger.info("Split command stdout: %s", result.stdout)
+    logger.info("Split command stderr: %s", result.stderr)
 
     if result.returncode != 0:
-        stderr_str = result.stderr.decode("utf-8")
-        raise Exception(f"Error splitting the model: {stderr_str}")
-    print("Model split successfully!")
+        raise RuntimeError(f"Error splitting the model: {result.stderr}")
+    logger.info("Model split successfully!")
 
-    # remove the original model file if needed
     if os.path.exists(model_path):
         os.remove(model_path)
 
     model_file_prefix = model_path_prefix.split("/")[-1]
-    print(f"Model file name prefix: {model_file_prefix}")
+    logger.info("Model file name prefix: %s", model_file_prefix)
     sharded_model_files = [
         f
         for f in os.listdir(outdir)
         if f.startswith(model_file_prefix) and f.endswith(".gguf")
     ]
-    if sharded_model_files:
-        print(f"Sharded model files: {sharded_model_files}")
-        if export_to_org and org_token != "":
-            api = HfApi(token=org_token)
+
+    if not sharded_model_files:
+        raise RuntimeError("No sharded files found.")
+
+    logger.info("Sharded model files: %s", sharded_model_files)
+    api = HfApi(token=org_token if (export_to_org and org_token) else oauth_token.token)
+
+    for file in sharded_model_files:
+        file_path = os.path.join(outdir, file)
+        logger.info("Uploading file: %s", file_path)
+        try:
+            api.upload_file(
+                path_or_fileobj=file_path,
+                path_in_repo=file,
+                repo_id=repo_id,
+            )
+        except Exception as e:
+            raise RuntimeError(f"Error uploading file {file_path}: {e}") from e
+
+    logger.info("Sharded model has been uploaded successfully!")
+
+
+def get_new_model_card(
+    original_card: ModelCard,
+    original_model_id: str,
+    gguf_files: list,
+    new_repo_url: str,
+    split_model: bool,
+) -> ModelCard:
+    version = get_llama_cpp_version()
+    model_card = original_card.copy()
+    model_card.data.tags = (model_card.data.tags or []) + [
+        "antigma",
+        "quantize-my-repo",
+    ]
+
+    # Format the table rows
+    table_rows = []
+    for file_info in gguf_files:
+        name, _, size, method = file_info
+        if split_model:
+            display_name = name[:-5]
         else:
-            api = HfApi(token=oauth_token.token)
-        for file in sharded_model_files:
-            file_path = os.path.join(outdir, file)
-            print(f"Uploading file: {file_path}")
-            try:
-                api.upload_file(
-                    path_or_fileobj=file_path,
-                    path_in_repo=file,
-                    repo_id=repo_id,
-                )
-            except Exception as e:
-                raise Exception(f"Error uploading file {file_path}: {e}")
-    else:
-        raise Exception("No sharded files found.")
+            display_name = f"[{name}]({new_repo_url}/blob/main/{name})"
+        table_rows.append(f"{display_name}|{method}|{size:.2f} GB|{split_model}")
+
+    model_card.text = f"""
+*Produced by [Antigma Labs](https://antigma.ai), [Antigma Quantize Space](https://huggingface.co/spaces/Antigma/quantize-my-repo)*
+
+*Follow Antigma Labs in X [https://x.com/antigma_labs](https://x.com/antigma_labs)*
+
+*Antigma's GitHub Homepage [https://github.com/AntigmaLabs](https://github.com/AntigmaLabs)*
+
+## Quantization Format (GGUF)
+We use <a href="https://github.com/ggml-org/llama.cpp">llama.cpp</a> release <a href="https://github.com/ggml-org/llama.cpp/releases/tag/{version}">{version}</a> for quantization.
+Original model: https://huggingface.co/{original_model_id}
+
+## Download a file (not the whole branch) from below:
+| Filename | Quant type | File Size | Split |
+| -------- | ---------- | --------- | ----- |
+| {'|'.join(table_rows)}
+
+## Original Model Card
+{original_card.text}
+
+## Downloading using huggingface-cli
+<details>
+  <summary>Click to view download instructions</summary>
+First, make sure you have hugginface-cli installed:
 
-    print("Sharded model has been uploaded successfully!")
+```
+pip install -U "huggingface_hub[cli]"
+```
+
+Then, you can target the specific file you want:
+
+```
+huggingface-cli download {new_repo_url} --include "{gguf_files[0][0]}" --local-dir ./
+```
+
+If the model is bigger than 50GB, it will have been split into multiple files. In order to download them all to a local folder, run:
+
+```
+huggingface-cli download {new_repo_url} --include "{gguf_files[0][0]}/*" --local-dir ./
+```
+
+You can either specify a new local-dir (e.g. deepseek-ai_DeepSeek-V3-0324-Q8_0) or it will be in default hugging face cache
+
+</details>
+"""
+    return model_card
 
 
 def process_model(
-    model_id,
-    q_method,
-    use_imatrix,
-    imatrix_q_method,
-    private_repo,
-    train_data_file,
-    split_model,
-    split_max_tensors,
-    split_max_size,
-    export_to_org,
-    repo_owner,
-    org_token,
+    model_id: str,
+    q_method: str | list,
+    use_imatrix: bool,
+    imatrix_q_method: str,
+    private_repo: bool,
+    train_data_file: gr.File | None,
+    split_model: bool,
+    split_max_tensors: int,
+    split_max_size: str | None,
+    export_to_org: bool,
+    repo_owner: str,
+    org_token: str | None,
     oauth_token: gr.OAuthToken | None,
-):
+) -> tuple[str, str]:
     if oauth_token is None or oauth_token.token is None:
         raise gr.Error("You must be logged in to use quantize-my-repo")
     try:
         whoami(oauth_token.token)
     except Exception as e:
-        raise gr.Error("You must be logged in to use quantize-my-repo")
+        raise gr.Error("You must be logged in to use quantize-my-repo") from e
 
     user_info = whoami(oauth_token.token)
     username = user_info["name"]
@@ -277,15 +289,17 @@ def process_model(
 
     current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     logger.info(
-        f"Time {current_time}, Username {username}, Model_ID, {model_id}, q_method {','.join(q_method)}"
+        "Time %s, Username %s, Model_ID %s, q_method %s",
+        current_time,
+        username,
+        model_id,
+        ",".join(q_method) if isinstance(q_method, list) else q_method,
     )
 
     repo_namespace = get_repo_namespace(repo_owner, username, user_orgs)
     model_name = model_id.split("/")[-1]
     try:
-        api_token = (
-            org_token if (export_to_org and org_token != "") else oauth_token.token
-        )
+        api_token = org_token if (export_to_org and org_token) else oauth_token.token
         api = HfApi(token=api_token)
 
         dl_pattern = ["*.md", "*.json", "*.model"]
@@ -297,7 +311,7 @@ def process_model(
             )
             else "*.bin"
         )
-        dl_pattern += [pattern]
+        dl_pattern.append(pattern)
 
         os.makedirs(downloads_dir, exist_ok=True)
         os.makedirs(outputs_dir, exist_ok=True)
@@ -306,10 +320,7 @@ def process_model(
             fp16 = str(Path(outdir) / f"{model_name}.fp16.gguf")
 
             with tempfile.TemporaryDirectory(dir=downloads_dir) as tmpdir:
-                print(datetime.now().strftime("%Y-%m-%d %H:%M:%S") + " Start download")
-                logger.info(
-                    datetime.now().strftime("%Y-%m-%d %H:%M:%S") + " Start download"
-                )
+                logger.info("Start download")
                 local_dir = Path(tmpdir) / model_name
                 api.snapshot_download(
                     repo_id=model_id,
@@ -323,17 +334,10 @@ def process_model(
                 if os.path.exists(adapter_config_dir) and not os.path.exists(
                     config_dir
                 ):
-                    raise Exception(
+                    raise RuntimeError(
                         "adapter_config.json is present. If converting LoRA, use GGUF-my-lora."
                     )
-                print(
-                    datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-                    + " Download successfully"
-                )
-                logger.info(
-                    datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-                    + " Download successfully"
-                )
+                logger.info("Download successfully")
 
                 result = subprocess.run(
                     [
@@ -348,15 +352,10 @@ def process_model(
                     shell=False,
                     capture_output=True,
                 )
-                print(
-                    datetime.now().strftime("%Y-%m-%d %H:%M:%S") + " Converted to f16"
-                )
-                logger.info(
-                    datetime.now().strftime("%Y-%m-%d %H:%M:%S") + " Converted to f16"
-                )
+                logger.info("Converted to f16")
 
                 if result.returncode != 0:
-                    raise Exception(
+                    raise RuntimeError(
                         f"Error converting to fp16: {result.stderr.decode()}"
                     )
             shutil.rmtree(downloads_dir)
@@ -369,7 +368,9 @@ def process_model(
                     else "llama.cpp/groups_merged.txt"
                 )
                 if not os.path.isfile(train_data_path):
-                    raise Exception(f"Training data not found: {train_data_path}")
+                    raise FileNotFoundError(
+                        f"Training data not found: {train_data_path}"
+                    )
                 generate_importance_matrix(fp16, train_data_path, imatrix_path)
 
             quant_methods = (
@@ -381,11 +382,7 @@ def process_model(
 
             gguf_files = []
             for method in quant_methods:
-                print(datetime.now().strftime("%Y-%m-%d %H:%M:%S") + " Begin quantize")
-                logger.info(
-                    datetime.now().strftime("%Y-%m-%d %H:%M:%S") + " Begin quantize"
-                )
-
+                logger.info("Begin quantize")
                 name = (
                     f"{model_name.lower()}-{method.lower()}-{suffix}.gguf"
                     if suffix
@@ -406,18 +403,13 @@ def process_model(
                 )
                 result = subprocess.run(quant_cmd, shell=False, capture_output=True)
                 if result.returncode != 0:
-                    raise Exception(
+                    raise RuntimeError(
                         f"Quantization failed ({method}): {result.stderr.decode()}"
                     )
                 size = os.path.getsize(path) / 1024 / 1024 / 1024
                 gguf_files.append((name, path, size, method))
 
-            print(
-                datetime.now().strftime("%Y-%m-%d %H:%M:%S") + " Quantize successfully!"
-            )
-            logger.info(
-                datetime.now().strftime("%Y-%m-%d %H:%M:%S") + " Quantize successfully!"
-            )
+            logger.info("Quantize successfully!")
 
             suffix_for_repo = (
                 f"{imatrix_q_method}-imat" if use_imatrix else "-".join(quant_methods)
@@ -428,13 +420,12 @@ def process_model(
             )
 
             try:
-                card = ModelCard.load(model_id, token=oauth_token.token)
-            except:
-                card = ModelCard("")
-            card.data.tags = (card.data.tags or []) + ["antigma", "quantize-my-repo"]
-            card.data.base_model = model_id
-            card.text = dedent(
-                get_llama_cpp_notes(gguf_files, new_repo_url, split_model, model_id)
+                original_card = ModelCard.load(model_id, token=oauth_token.token)
+            except Exception:
+                original_card = ModelCard("")
+
+            card = get_new_model_card(
+                original_card, model_id, gguf_files, new_repo_url, split_model
             )
             readme_path = Path(outdir) / "README.md"
             card.save(readme_path)
@@ -478,6 +469,7 @@ def process_model(
 css = """/* Custom CSS to allow scrolling */
 .gradio-container {overflow-y: auto;}
 """
+
 model_id = HuggingfaceHubSearch(
     label="Hub Model ID",
     placeholder="Search for model id on Huggingface",
@@ -578,6 +570,7 @@ iface = gr.Interface(
     description="We take your Hugging Face repo — a terrific repo — we quantize it, we package it beautifully, and we give you your very own repo. It's smart. It's efficient. It's huge. You're gonna love it.",
     api_name=False,
 )
+
 with gr.Blocks(css=".gradio-container {overflow-y: auto;}") as demo:
     gr.Markdown("Logged in, you must be. Classy, secure, and victorious, it keeps us.")
     gr.LoginButton(min_width=250)