Update app.py

app.py

@@ -1,10 +1,10 @@
-import spaces
 import os
 import re
 import json
 import tempfile
 import zipfile
 import gradio as gr
+import spaces
 from huggingface_hub import hf_hub_download
 
 # ---- LLM: llama.cpp via llama_cpp_agent ----
@@ -15,11 +15,9 @@ from llama_cpp_agent.providers import LlamaCppPythonProvider
 # ----------------------
 # Model configuration
 # ----------------------
-# You can change these defaults in the UI
 DEFAULT_REPO_ID = "tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF"
 DEFAULT_FILENAME = "neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf"
 MODELS_DIR = "models"
-
 os.makedirs(MODELS_DIR, exist_ok=True)
 
 def ensure_model(repo_id: str, filename: str) -> str:
@@ -28,13 +26,16 @@ def ensure_model(repo_id: str, filename: str) -> str:
         hf_hub_download(repo_id=repo_id, filename=filename, local_dir=MODELS_DIR)
     return local_path
 
-def build_agent(model_path: str, n_ctx=8192, n_gpu_layers=81, n_batch=1024, flash_attn=True):
+# GPU context ensures model loads only when a GPU session is active
+@spaces.GPU(duration=120)
+def package_with_llm(ai_text, repo_id, filename, temperature, top_p, top_k, repeat_penalty, max_tokens, force_heuristic):
+    model_path = ensure_model(repo_id, filename)
     llm = Llama(
         model_path=model_path,
-        n_ctx=n_ctx,
-        n_gpu_layers=n_gpu_layers,
-        n_batch=n_batch,
-        flash_attn=flash_attn,
+        n_ctx=8192,
+        n_gpu_layers=81,
+        n_batch=1024,
+        flash_attn=True,
     )
     provider = LlamaCppPythonProvider(llm)
     agent = LlamaCppAgent(
@@ -52,21 +53,16 @@ def build_agent(model_path: str, n_ctx=8192, n_gpu_layers=81, n_batch=1024, flas
         predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
         debug_output=False
     )
-    return agent, provider
 
-JSON_FALLBACK_NAME = "project.txt"
-    
-@spaces.GPU(duration=120) 
-def call_llm_manifest(agent, provider, text, temperature=0.2, top_p=0.9, top_k=40, repeat_penalty=1.1, max_tokens=2048):
-    # Instruction prompt. Model must respond with STRICT JSON.
     prompt = f"""
-Read the following AI project description and return ONLY JSON. 
+Read the following AI project description and return ONLY JSON.
 Output schema (strict):
 [{{"filename": "server.js", "content": "// code..."}}]
 
 AI project description:
-{text}
+{ai_text}
 """
+
     settings = provider.get_provider_default_settings()
     settings.temperature = float(temperature)
     settings.top_p = float(top_p)
@@ -76,38 +72,32 @@ AI project description:
     settings.stream = False
 
     out = agent.get_chat_response(prompt, llm_sampling_settings=settings, print_output=False)
-    # Try to extract a JSON array from the output robustly
-    json_text = None
+
+    # Parse JSON output robustly
     try:
-        # Prefer the largest bracketed array slice
-        start = out.find('[')
-        end = out.rfind(']')
-        if start != -1 and end != -1 and end > start:
-            json_text = out[start:end+1]
-            manifest = json.loads(json_text)
+        start, end = out.find('['), out.rfind(']')
+        if start != -1 and end > start:
+            manifest = json.loads(out[start:end+1])
         else:
-            raise ValueError("No JSON array found")
+            raise ValueError("No JSON found")
     except Exception:
-        # Fallback: single-file package of raw output for transparency
-        manifest = [ {"filename": JSON_FALLBACK_NAME, "content": out} ]
-    return manifest
+        manifest = [{"filename": "project.txt", "content": out}]
+
+    # Heuristic fallback if model fails
+    if force_heuristic or (len(manifest) == 1 and manifest[0]["filename"] == "project.txt"):
+        manifest = naive_regex_merge(ai_text)
+
+    return create_zip_from_manifest(manifest)
 
 def naive_regex_merge(text):
-    """
-    Heuristic backup that maps code fences to probable filenames by scanning nearby lines.
-    This runs only when the model output is a single fallback file OR user ticks 'Force Heuristic Merge'.
-    """
     blocks = []
-    # Find all triple-backtick code blocks
     code_pattern = re.compile(r"```([a-zA-Z0-9]*)\n(.*?)```", re.DOTALL)
-    # Find filename candidates in preceding lines such as '### STEP: server.js' or '`server.js`'
     lines = text.splitlines()
     candidates = []
-    for i, line in enumerate(lines):
-        m = re.search(r"([A-Za-z0-9_\\-./]+?\\.[A-Za-z0-9]+)", line)
+    for line in lines:
+        m = re.search(r"([A-Za-z0-9_\-./]+?\.[A-Za-z0-9]+)", line)
         if m:
             candidates.append(m.group(1))
-
     for idx, m in enumerate(code_pattern.finditer(text)):
         lang = m.group(1) or "txt"
         code = m.group(2)
@@ -120,7 +110,7 @@ def create_zip_from_manifest(manifest):
     zip_path = os.path.join(temp_dir, "project.zip")
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
         for item in manifest:
-            fname = item.get("filename", JSON_FALLBACK_NAME).lstrip("/")
+            fname = item.get("filename", "project.txt").lstrip("/")
             content = item.get("content", "")
             fpath = os.path.join(temp_dir, fname)
             os.makedirs(os.path.dirname(fpath), exist_ok=True)
@@ -129,28 +119,12 @@ def create_zip_from_manifest(manifest):
             z.write(fpath, arcname=fname)
     return zip_path
 
-def package_with_llm(ai_text, repo_id, filename, temperature, top_p, top_k, repeat_penalty, max_tokens, force_heuristic):
-    model_path = ensure_model(repo_id, filename)
-    agent, provider = build_agent(model_path=model_path)
-
-    manifest = call_llm_manifest(
-        agent, provider, ai_text,
-        temperature=temperature, top_p=top_p, top_k=top_k,
-        repeat_penalty=repeat_penalty, max_tokens=max_tokens
-    )
-
-    # If model failed to JSON-ify properly (single fallback) or user forces merge, try heuristic merge
-    if force_heuristic or (len(manifest) == 1 and manifest[0]["filename"] == JSON_FALLBACK_NAME):
-        heuristic = naive_regex_merge(ai_text)
-        if heuristic:
-            manifest = heuristic
-
-    zip_path = create_zip_from_manifest(manifest)
-    return zip_path
-
-with gr.Blocks(title="AI Project Packager (GGUF, llama.cpp)") as demo:
-    gr.Markdown("# AI Project Packager (GGUF, llama.cpp)")
-    gr.Markdown("Paste an AI-generated multi-file project description. A local GGUF model will infer filenames and contents, then return a downloadable ZIP.")
+# --------------------------
+# Gradio Interface
+# --------------------------
+with gr.Blocks(title="AI Project Packager (GGUF, ZeroGPU)") as demo:
+    gr.Markdown("# AI Project Packager (GGUF, ZeroGPU)")
+    gr.Markdown("Uses ephemeral GPU power to infer filenames & package AI-generated code projects.")
 
     with gr.Row():
         ai_text = gr.Textbox(lines=24, label="Paste AI response here")