Save the model to the data directory

app.py

@@ -15,45 +15,29 @@ import numpy as np
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
 CONVERSION_SCRIPT = "./llama.cpp/convert_hf_to_gguf.py"
 
-# Get Hugging Face token from environment variable
-HF_TOKEN = os.environ.get("HF_TOKEN")
-
-# Set up persistent storage paths
 log_dir = "/data/logs"
 downloads_dir = "/data/downloads"
 outputs_dir = "/data/outputs"
-models_dir = "/data/models"
-
-# Create directories if they don't exist
 os.makedirs(log_dir, exist_ok=True)
-os.makedirs(downloads_dir, exist_ok=True)
-os.makedirs(outputs_dir, exist_ok=True)
-os.makedirs(models_dir, exist_ok=True)
 
 logging.basicConfig(
     filename=os.path.join(log_dir, "app.log"),
     level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s",
+    format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
 logger = logging.getLogger(__name__)
 
-
-def get_llama_cpp_notes(
-    gguf_files,
-    new_repo_url,
-    split_model,
-    model_id=None,
-):
+def get_llama_cpp_notes(gguf_files, new_repo_url, split_model, model_id = None,):
     try:
         result = subprocess.run(
-            ["git", "-C", "./llama.cpp", "describe", "--tags", "--always"],
+            ['git', '-C', './llama.cpp', 'describe', '--tags', '--always'],
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             check=True,
-            text=True,
+            text=True
         )
-        version = result.stdout.strip().split("-")[0]
+        version = result.stdout.strip().split('-')[0]
         text = f"""
 *Produced by [Antigma Labs](https://antigma.ai)*
 ## llama.cpp quantization
@@ -62,8 +46,7 @@ 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
 ```
-{{system_prompt}}
-{{prompt}}
+<｜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 |
@@ -95,51 +78,32 @@ You can either specify a new local-dir (deepseek-ai_DeepSeek-V3-0324-Q8_0) or do
 
 
 def get_repo_namespace(repo_owner, username, user_orgs):
-    if repo_owner == "self":
+    if repo_owner == 'self':
         return username
     for org in user_orgs:
-        if org["name"] == repo_owner:
-            return org["name"]
+        if org['name'] == repo_owner:
+            return org['name']
     raise ValueError(f"Invalid repo_owner: {repo_owner}")
 
-
 def escape(s: str) -> str:
-    return (
-        s.replace("&", "&amp;")
-        .replace("<", "&lt;")
-        .replace(">", "&gt;")
-        .replace('"', "&quot;")
-        .replace("\n", "<br/>")
-    )
-
+    return s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;").replace("\n", "<br/>")
 
 def toggle_repo_owner(export_to_org, oauth_token: gr.OAuthToken | None):
     if oauth_token is None or oauth_token.token is None:
         raise gr.Error("You must be logged in to use GGUF-my-repo")
     if not export_to_org:
-        return gr.update(visible=False, choices=["self"], value="self"), gr.update(
-            visible=False, value=""
-        )
+        return gr.update(visible=False, choices=["self"], value="self"), gr.update(visible=False, value="")
     info = whoami(oauth_token.token)
     orgs = [org["name"] for org in info.get("orgs", [])]
-    return gr.update(visible=True, choices=["self"] + orgs, value="self"), gr.update(
-        visible=True
-    )
-
-
+    return gr.update(visible=True, choices=["self"] + orgs, value="self"), gr.update(visible=True)
 def generate_importance_matrix(model_path: str, train_data_path: str, output_path: str):
     imatrix_command = [
         "./llama.cpp/llama-imatrix",
-        "-m",
-        model_path,
-        "-f",
-        train_data_path,
-        "-ngl",
-        "99",
-        "--output-frequency",
-        "10",
-        "-o",
-        output_path,
+        "-m", model_path,
+        "-f", train_data_path,
+        "-ngl", "99",
+        "--output-frequency", "10",
+        "-o", output_path,
     ]
 
     if not os.path.isfile(model_path):
@@ -151,9 +115,7 @@ def generate_importance_matrix(model_path: str, train_data_path: str, output_pat
     try:
         process.wait(timeout=60)  # added wait
     except subprocess.TimeoutExpired:
-        print(
-            "Imatrix computation timed out. Sending SIGINT to allow graceful termination..."
-        )
+        print("Imatrix computation timed out. Sending SIGINT to allow graceful termination...")
         process.send_signal(signal.SIGINT)
         try:
             process.wait(timeout=5)  # grace period
@@ -163,17 +125,7 @@ def generate_importance_matrix(model_path: str, train_data_path: str, output_pat
 
     print("Importance matrix generation completed.")
 
-
-def split_upload_model(
-    model_path: str,
-    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,
-):
+def split_upload_model(model_path: str, 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}")
 
@@ -192,9 +144,7 @@ def split_upload_model(
         split_cmd.append(str(split_max_tensors))
 
     # args for output
-    model_path_prefix = ".".join(
-        model_path.split(".")[:-1]
-    )  # remove the file extension
+    model_path_prefix = '.'.join(model_path.split('.')[:-1]) # remove the file extension
     split_cmd.append(model_path)
     split_cmd.append(model_path_prefix)
 
@@ -213,19 +163,15 @@ def split_upload_model(
     if os.path.exists(model_path):
         os.remove(model_path)
 
-    model_file_prefix = model_path_prefix.split("/")[-1]
+    model_file_prefix = model_path_prefix.split('/')[-1]
     print(f"Model file name prefix: {model_file_prefix}")
-    sharded_model_files = [
-        f
-        for f in os.listdir(outdir)
-        if f.startswith(model_file_prefix) and f.endswith(".gguf")
-    ]
+    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 export_to_org and org_token!="":
+          api = HfApi(token = org_token)
         else:
-            api = HfApi(token=oauth_token.token)
+          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}")
@@ -242,22 +188,9 @@ def split_upload_model(
 
     print("Sharded model has been uploaded successfully!")
 
-
-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,
-    oauth_token: gr.OAuthToken | None,
-):
+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, oauth_token: gr.OAuthToken | None):
     if oauth_token is None or oauth_token.token is None:
         raise gr.Error("You must be logged in to use GGUF-my-repo")
 
@@ -267,176 +200,91 @@ def process_model(
     if not export_to_org:
         repo_owner = "self"
 
+
     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)}"
-    )
+    logger.info(f"Time {current_time}, Username {username}, Model_ID, {model_id}, q_method {','.join(q_method)}")
 
     repo_namespace = get_repo_namespace(repo_owner, username, user_orgs)
-    model_name = model_id.split("/")[-1]
+    model_name = model_id.split('/')[-1]
     try:
-        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"]
-        pattern = (
-            "*.safetensors"
-            if any(
-                f.path.endswith(".safetensors")
-                for f in api.list_repo_tree(repo_id=model_id, recursive=True)
-            )
-            else "*.bin"
-        )
-        dl_pattern += [pattern]
-
-        # Use persistent directories
-        model_download_dir = os.path.join(downloads_dir, model_name)
-        os.makedirs(model_download_dir, exist_ok=True)
-
-        with tempfile.TemporaryDirectory(dir=outputs_dir) as outdir:
-            fp16 = str(Path(outdir) / f"{model_name}.fp16.gguf")
-
-            with tempfile.TemporaryDirectory(dir=model_download_dir) as tmpdir:
-                local_dir = Path(tmpdir) / model_name
-                api.snapshot_download(
-                    repo_id=model_id,
-                    local_dir=local_dir,
-                    local_dir_use_symlinks=False,
-                    allow_patterns=dl_pattern,
-                )
-
-                config_dir = local_dir / "config.json"
-                adapter_config_dir = local_dir / "adapter_config.json"
-                if os.path.exists(adapter_config_dir) and not os.path.exists(
-                    config_dir
-                ):
-                    raise Exception(
-                        "adapter_config.json is present. If converting LoRA, use GGUF-my-lora."
-                    )
-
-                result = subprocess.run(
-                    [
-                        "python",
-                        CONVERSION_SCRIPT,
-                        local_dir,
-                        "--outtype",
-                        "f16",
-                        "--outfile",
-                        fp16,
-                    ],
-                    shell=False,
-                    capture_output=True,
-                )
-                if result.returncode != 0:
-                    raise Exception(
-                        f"Error converting to fp16: {result.stderr.decode()}"
-                    )
-
-            imatrix_path = Path(outdir) / "imatrix.dat"
-            if use_imatrix:
-                train_data_path = (
-                    train_data_file.name
-                    if train_data_file
-                    else "llama.cpp/groups_merged.txt"
-                )
-                if not os.path.isfile(train_data_path):
-                    raise Exception(f"Training data not found: {train_data_path}")
-                generate_importance_matrix(fp16, train_data_path, imatrix_path)
-
-            quant_methods = (
-                [imatrix_q_method]
-                if use_imatrix
-                else (q_method if isinstance(q_method, list) else [q_method])
-            )
-            suffix = "imat" if use_imatrix else None
-
-            gguf_files = []
-            for method in quant_methods:
-                name = (
-                    f"{model_name.lower()}-{method.lower()}-{suffix}.gguf"
-                    if suffix
-                    else f"{model_name.lower()}-{method.lower()}.gguf"
-                )
-                path = str(Path(outdir) / name)
-                quant_cmd = (
-                    [
-                        "./llama.cpp/llama-quantize",
-                        "--imatrix",
-                        imatrix_path,
-                        fp16,
-                        path,
-                        method,
-                    ]
-                    if use_imatrix
-                    else ["./llama.cpp/llama-quantize", fp16, path, method]
-                )
-                result = subprocess.run(quant_cmd, shell=False, capture_output=True)
-                if result.returncode != 0:
-                    raise Exception(
-                        f"Quantization failed ({method}): {result.stderr.decode()}"
-                    )
-                size = os.path.getsize(path) / 1024 / 1024 / 1024
-                gguf_files.append((name, path, size, method))
-
-            suffix_for_repo = (
-                f"{imatrix_q_method}-imat" if use_imatrix else "-".join(quant_methods)
-            )
-            repo_id = f"{repo_namespace}/{model_name}-{suffix_for_repo}-GGUF"
-            new_repo_url = api.create_repo(
-                repo_id=repo_id, exist_ok=True, private=private_repo
-            )
-
-            try:
-                card = ModelCard.load(model_id, token=oauth_token.token)
-            except:
-                card = ModelCard("")
-            card.data.tags = (card.data.tags or []) + ["llama-cpp", "gguf-my-repo"]
-            card.data.base_model = model_id
-            card.text = dedent(
-                get_llama_cpp_notes(gguf_files, new_repo_url, split_model, model_id)
-            )
-            readme_path = Path(outdir) / "README.md"
-            card.save(readme_path)
-            for name, path, _, _ in gguf_files:
-                if split_model:
-                    split_upload_model(
-                        path,
-                        outdir,
-                        repo_id,
-                        oauth_token,
-                        split_max_tensors,
-                        split_max_size,
-                        org_token,
-                        export_to_org,
-                    )
-                else:
-                    api.upload_file(
-                        path_or_fileobj=path, path_in_repo=name, repo_id=repo_id
-                    )
-            if use_imatrix and os.path.isfile(imatrix_path):
-                api.upload_file(
-                    path_or_fileobj=imatrix_path,
-                    path_in_repo="imatrix.dat",
-                    repo_id=repo_id,
-                )
-            api.upload_file(
-                path_or_fileobj=readme_path, path_in_repo="README.md", repo_id=repo_id
-            )
-
-            return (
-                f'<h1>✅ DONE</h1><br/>Repo: <a href="{new_repo_url}" target="_blank" style="text-decoration:underline">{repo_id}</a>',
-                f"llama{np.random.randint(9)}.png",
-            )
+      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"]
+      pattern = "*.safetensors" if any(
+          f.path.endswith(".safetensors")
+          for f in api.list_repo_tree(repo_id=model_id, recursive=True)
+      ) else "*.bin"
+      dl_pattern += [pattern]
+
+      os.makedirs(downloads_dir, exist_ok=True)
+      os.makedirs(outputs_dir, exist_ok=True)
+
+      with tempfile.TemporaryDirectory(dir=outputs_dir) as outdir:
+          fp16 = str(Path(outdir)/f"{model_name}.fp16.gguf")
+
+          with tempfile.TemporaryDirectory(dir=downloads_dir) as tmpdir:
+              local_dir = Path(tmpdir)/model_name
+              api.snapshot_download(repo_id=model_id, local_dir=local_dir, local_dir_use_symlinks=False, allow_patterns=dl_pattern)
+
+              config_dir = local_dir/"config.json"
+              adapter_config_dir = local_dir/"adapter_config.json"
+              if os.path.exists(adapter_config_dir) and not os.path.exists(config_dir):
+                  raise Exception("adapter_config.json is present. If converting LoRA, use GGUF-my-lora.")
+
+              result = subprocess.run(["python", CONVERSION_SCRIPT, local_dir, "--outtype", "f16", "--outfile", fp16], shell=False, capture_output=True)
+              if result.returncode != 0:
+                  raise Exception(f"Error converting to fp16: {result.stderr.decode()}")
+
+          imatrix_path = Path(outdir)/"imatrix.dat"
+          if use_imatrix:
+              train_data_path = train_data_file.name if train_data_file else "llama.cpp/groups_merged.txt"
+              if not os.path.isfile(train_data_path):
+                  raise Exception(f"Training data not found: {train_data_path}")
+              generate_importance_matrix(fp16, train_data_path, imatrix_path)
+
+          quant_methods = [imatrix_q_method] if use_imatrix else (q_method if isinstance(q_method, list) else [q_method])
+          suffix = "imat" if use_imatrix else None
+
+          gguf_files = []
+          for method in quant_methods:
+              name = f"{model_name.lower()}-{method.lower()}-{suffix}.gguf" if suffix else f"{model_name.lower()}-{method.lower()}.gguf"
+              path = str(Path(outdir)/name)
+              quant_cmd = ["./llama.cpp/llama-quantize", "--imatrix", imatrix_path, fp16, path, method] if use_imatrix else ["./llama.cpp/llama-quantize", fp16, path, method]
+              result = subprocess.run(quant_cmd, shell=False, capture_output=True)
+              if result.returncode != 0:
+                  raise Exception(f"Quantization failed ({method}): {result.stderr.decode()}")
+              size = os.path.getsize(path)/1024/1024/1024
+              gguf_files.append((name, path, size, method))
+
+          suffix_for_repo = f"{imatrix_q_method}-imat" if use_imatrix else "-".join(quant_methods)
+          repo_id = f"{repo_namespace}/{model_name}-{suffix_for_repo}-GGUF"
+          new_repo_url = api.create_repo(repo_id=repo_id, exist_ok=True, private=private_repo)
+
+          try:
+              card = ModelCard.load(model_id, token=oauth_token.token)
+          except:
+              card = ModelCard("")
+          card.data.tags = (card.data.tags or []) + ["llama-cpp", "gguf-my-repo"]
+          card.data.base_model = model_id
+          card.text = dedent(get_llama_cpp_notes(gguf_files, new_repo_url, split_model, model_id))
+          readme_path = Path(outdir)/"README.md"
+          card.save(readme_path)
+          for name, path, _, _ in gguf_files:
+              if split_model:
+                  split_upload_model(path, outdir, repo_id, oauth_token, split_max_tensors, split_max_size, org_token, export_to_org)
+              else:
+                  api.upload_file(path_or_fileobj=path, path_in_repo=name, repo_id=repo_id)
+          if use_imatrix and os.path.isfile(imatrix_path):
+              api.upload_file(path_or_fileobj=imatrix_path, path_in_repo="imatrix.dat", repo_id=repo_id)
+          api.upload_file(path_or_fileobj=readme_path, path_in_repo="README.md", repo_id=repo_id)
+
+          return (f'<h1>✅ DONE</h1><br/>Repo: <a href="{new_repo_url}" target="_blank" style="text-decoration:underline">{repo_id}</a>', f"llama{np.random.randint(9)}.png")
     except Exception as e:
-        raise (
-            f'<h1>❌ ERROR</h1><br/><pre style="white-space:pre-wrap;">{escape(str(e))}</pre>',
-            "error.png",
-        )
+        raise (f'<h1>❌ ERROR</h1><br/><pre style="white-space:pre-wrap;">{escape(str(e))}</pre>', "error.png")
 
 
-css = """/* Custom CSS to allow scrolling */
+css="""/* Custom CSS to allow scrolling */
 .gradio-container {overflow-y: auto;}
 """
 model_id = HuggingfaceHubSearch(
@@ -448,36 +296,30 @@ model_id = HuggingfaceHubSearch(
 export_to_org = gr.Checkbox(
     label="Export to Organization Repository",
     value=False,
-    info="If checked, you can select an organization to export to.",
+    info="If checked, you can select an organization to export to."
 )
 
 repo_owner = gr.Dropdown(
-    choices=["self"], value="self", label="Repository Owner", visible=False
+    choices=["self"],
+    value="self",
+    label="Repository Owner",
+    visible=False
 )
 
-org_token = gr.Textbox(label="Org Access Token", type="password", visible=False)
+org_token = gr.Textbox(
+    label="Org Access Token",
+    type="password",
+    visible=False
+)
 
 q_method = gr.Dropdown(
-    [
-        "Q2_K",
-        "Q3_K_S",
-        "Q3_K_M",
-        "Q3_K_L",
-        "Q4_0",
-        "Q4_K_S",
-        "Q4_K_M",
-        "Q5_0",
-        "Q5_K_S",
-        "Q5_K_M",
-        "Q6_K",
-        "Q8_0",
-    ],
+    ["Q2_K", "Q3_K_S", "Q3_K_M", "Q3_K_L", "Q4_0", "Q4_K_S", "Q4_K_M", "Q5_0", "Q5_K_S", "Q5_K_M", "Q6_K", "Q8_0"],
     label="Quantization Method",
     info="GGML quantization type",
     value="Q4_K_M",
     filterable=False,
     visible=True,
-    multiselect=True,
+    multiselect=True
 )
 
 imatrix_q_method = gr.Dropdown(
@@ -486,36 +328,44 @@ imatrix_q_method = gr.Dropdown(
     info="GGML imatrix quants type",
     value="IQ4_NL",
     filterable=False,
-    visible=False,
+    visible=False
 )
 
 use_imatrix = gr.Checkbox(
     value=False,
     label="Use Imatrix Quantization",
-    info="Use importance matrix for quantization.",
+    info="Use importance matrix for quantization."
 )
 
 private_repo = gr.Checkbox(
-    value=False, label="Private Repo", info="Create a private repo under your username."
+    value=False,
+    label="Private Repo",
+    info="Create a private repo under your username."
 )
 
-train_data_file = gr.File(label="Training Data File", file_types=["txt"], visible=False)
+train_data_file = gr.File(
+    label="Training Data File",
+    file_types=["txt"],
+    visible=False
+)
 
 split_model = gr.Checkbox(
-    value=False, label="Split Model", info="Shard the model using gguf-split."
+    value=False,
+    label="Split Model",
+    info="Shard the model using gguf-split."
 )
 
 split_max_tensors = gr.Number(
     value=256,
     label="Max Tensors per File",
     info="Maximum number of tensors per file when splitting model.",
-    visible=False,
+    visible=False
 )
 
 split_max_size = gr.Textbox(
     label="Max File Size",
     info="Maximum file size when splitting model (--split-max-size). May leave empty to use the default. Accepted suffixes: M, G. Example: 256M, 5G",
-    visible=False,
+    visible=False
 )
 
 iface = gr.Interface(
@@ -532,47 +382,35 @@ iface = gr.Interface(
         split_max_size,
         export_to_org,
         repo_owner,
-        org_token,
+        org_token
+    ],
+    outputs=[
+        gr.Markdown(label="Output"),
+        gr.Image(show_label=False)
     ],
-    outputs=[gr.Markdown(label="Output"), gr.Image(show_label=False)],
     title="Make your own GGUF Quants — faster than ever before, believe me.",
     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,
+    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)
 
-    export_to_org.change(
-        fn=toggle_repo_owner, inputs=[export_to_org], outputs=[repo_owner, org_token]
-    )
-
-    split_model.change(
-        fn=lambda sm: (gr.update(visible=sm), gr.update(visible=sm)),
-        inputs=split_model,
-        outputs=[split_max_tensors, split_max_size],
-    )
-    use_imatrix.change(
-        fn=lambda use: (
-            gr.update(visible=not use),
-            gr.update(visible=use),
-            gr.update(visible=use),
-        ),
-        inputs=use_imatrix,
-        outputs=[q_method, imatrix_q_method, train_data_file],
-    )
+
+
+    export_to_org.change(fn=toggle_repo_owner, inputs=[export_to_org], outputs=[repo_owner, org_token])
+
+    split_model.change(fn=lambda sm: (gr.update(visible=sm), gr.update(visible=sm)), inputs=split_model, outputs=[split_max_tensors, split_max_size])
+    use_imatrix.change(fn=lambda use: (gr.update(visible=not use), gr.update(visible=use), gr.update(visible=use)), inputs=use_imatrix, outputs=[q_method, imatrix_q_method, train_data_file])
 
     iface.render()
 
 
 def restart_space():
-    HfApi().restart_space(
-        repo_id="Antigma/quantize-my-repo", token=HF_TOKEN, factory_reboot=True
-    )
-
+    HfApi().restart_space(repo_id="Antigma/quantize-my-repo", token=HF_TOKEN, factory_reboot=True)
 
 scheduler = BackgroundScheduler()
 scheduler.add_job(restart_space, "interval", seconds=21600)
 scheduler.start()
 
-demo.queue(default_concurrency_limit=1, max_size=5).launch(debug=True, show_api=False)
+demo.queue(default_concurrency_limit=1, max_size=5).launch(debug=True, show_api=False)