Update app.py

app.py

@@ -1,28 +1,102 @@
-import spaces
-from model_runner import ModelRunner
-from parser import parse_to_files
-from zip_util import make_zip_from_files
 import gradio as gr
+import spaces
+
+# We delay all CUDA-related imports to inside the GPU-decorated function
+
+@spaces.GPU(duration=120)
+def process_code_generation(ai_response: str):
+    import json
+    import re
+    from huggingface_hub import hf_hub_download
+    from llama_cpp import Llama
+    from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
+    from llama_cpp_agent.providers import LlamaCppPythonProvider
+
+    # Download the model (cached if already exists)
+    model_path = hf_hub_download(
+        repo_id="tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF",
+        filename="neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf",
+        local_dir="./models"
+    )
+
+    # Initialize LLaMA model (CPU-safe)
+    llm = Llama(
+        model_path=model_path,
+        n_gpu_layers=0,     # IMPORTANT: ZeroGPU uses CPU under the hood
+        n_batch=512,
+        n_ctx=4096,
+        verbose=True
+    )
+    provider = LlamaCppPythonProvider(llm)
+    agent = LlamaCppAgent(
+        provider,
+        system_prompt="You are an AI code parser. Given a pasted AI response with multiple code files, extract the individual files and output a JSON object like {\"file.py\": \"...\", \"utils.py\": \"...\"}.",
+        predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
+        debug_output=False
+    )
+
+    # Prompt to interpret the response
+    user_prompt = f"""
+Given the following AI output that includes multiple code files, extract and convert it into a JSON object that maps filenames to code contents.
+
+Example output: 
+{{
+  "app.py": "print('Hello World')",
+  "utils/helper.py": "def greet(): return 'hi'"
+}}
 
-# Same model you referenced
-MODEL_REPO = "tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF"
-MODEL_FILE = "neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf"
+AI Output:
+{ai_response}
+"""
 
-runner = ModelRunner(repo_id=MODEL_REPO, filename=MODEL_FILE)
+    full_response = ""
+    for chunk in agent.get_chat_response(user_prompt, returns_streaming_generator=True):
+        full_response += chunk
+
+    # Try to parse as JSON
+    try:
+        parsed = json.loads(full_response)
+    except Exception:
+        # Fallback regex parsing if model output wasn't pure JSON
+        pattern = r"Filename:\s*(?P<fname>[\w\.\-/]+)\s*```(?:[a-zA-Z0-9]*)\n(?P<code>[\s\S]*?)```"
+        parsed = {}
+        for m in re.finditer(pattern, full_response):
+            parsed[m.group("fname").strip()] = m.group("code")
+
+        if not parsed:
+            parsed = {"output.txt": full_response}
+
+    # Write files and zip them
+    import os
+    import zipfile
+    import uuid
+
+    temp_dir = f"/tmp/code_files_{uuid.uuid4().hex}"
+    os.makedirs(temp_dir, exist_ok=True)
+
+    for filename, content in parsed.items():
+        file_path = os.path.join(temp_dir, filename)
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+
+    zip_path = f"/tmp/code_bundle_{uuid.uuid4().hex}.zip"
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(temp_dir):
+            for file in files:
+                full_path = os.path.join(root, file)
+                arcname = os.path.relpath(full_path, start=temp_dir)
+                zipf.write(full_path, arcname=arcname)
 
-@spaces.GPU(duration=120) 
-def process_and_zip(ai_response: str):
-    model_output = runner.interpret_code_description(ai_response)
-    file_list = parse_to_files(model_output)
-    zip_path = make_zip_from_files(file_list)
     return zip_path
 
+
 with gr.Blocks() as demo:
-    gr.Markdown("### Paste your AI output with multiple files below:")
-    input_box = gr.Textbox(lines=15, placeholder="Paste codegen or AI output here...")
-    submit = gr.Button("Generate Zip")
-    download = gr.File(label="Download zip")
+    gr.Markdown("### Paste your AI response below (with multiple files), and get a downloadable zip.")
+    input_box = gr.Textbox(lines=15, label="AI Response Input")
+    output_file = gr.File(label="Download Zip")
+    run_btn = gr.Button("Generate Zip")
 
-    submit.click(fn=process_and_zip, inputs=input_box, outputs=download)
+    run_btn.click(fn=process_code_generation, inputs=input_box, outputs=output_file)
 
 demo.launch()