File size: 2,959 Bytes
5da2830
548ac73
5da2830
548ac73
5da2830
548ac73
 
 
5da2830
 
 
8361a08
5da2830
 
 
548ac73
5da2830
8361a08
 
5da2830
 
c1a6fab
5da2830
 
 
 
 
 
 
 
 
 
 
 
 
 
c1a6fab
5da2830
 
c1a6fab
5da2830
 
 
 
 
 
 
c1a6fab
5da2830
 
 
 
 
548ac73
5da2830
 
 
 
 
 
 
c1a6fab
5da2830
548ac73
5da2830
 
 
 
 
 
 
 
 
 
c1a6fab
5da2830
 
 
c1a6fab
5da2830
 
 
 
 
 
 
 
 
c1a6fab
5da2830
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import spaces
import os
import gradio as gr
import zipfile
import tempfile
from llama_cpp import Llama
from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
from llama_cpp_agent.providers import LlamaCppPythonProvider
from llama_cpp_agent.chat_history import BasicChatHistory
from llama_cpp_agent.chat_history.messages import Roles
from huggingface_hub import hf_hub_download

MODEL_REPO = "tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF"
MODEL_FILENAME = "neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf"
MODEL_PATH = f"models/{MODEL_FILENAME}"

hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME, local_dir="./models")

@spaces.GPU(duration=120)
def generate_code_zip(prompt):
    chat_template = MessagesFormatterType.GEMMA_2

    llm = Llama(
        model_path=MODEL_PATH,
        flash_attn=True,
        n_gpu_layers=81,
        n_batch=1024,
        n_ctx=8192,
    )
    provider = LlamaCppPythonProvider(llm)
    agent = LlamaCppAgent(
        provider,
        system_prompt="You are an AI code generation agent. Respond with code blocks per file name.",
        predefined_messages_formatter_type=chat_template,
        debug_output=True
    )

    messages = BasicChatHistory()
    messages.add_message({"role": Roles.user, "content": prompt})

    stream = agent.get_chat_response(
        prompt,
        llm_sampling_settings=provider.get_provider_default_settings(),
        chat_history=messages,
        returns_streaming_generator=False,
        print_output=True
    )

    response = stream if isinstance(stream, str) else "".join(stream)
    file_map = parse_code_blocks(response)
    
    temp_dir = tempfile.mkdtemp()
    zip_path = os.path.join(temp_dir, "generated_code.zip")

    with zipfile.ZipFile(zip_path, "w") as zipf:
        for filename, code in file_map.items():
            full_path = os.path.join(temp_dir, filename)
            os.makedirs(os.path.dirname(full_path), exist_ok=True)
            with open(full_path, "w", encoding="utf-8") as f:
                f.write(code)
            zipf.write(full_path, arcname=filename)

    return zip_path

def parse_code_blocks(text):
    import re
    file_blocks = {}
    pattern = r"###\s*(.*?)\n```(?:[a-zA-Z0-9+]*\n)?(.*?)```"
    matches = re.findall(pattern, text, re.DOTALL)
    for filename, content in matches:
        cleaned_filename = filename.strip()
        cleaned_code = content.strip()
        file_blocks[cleaned_filename] = cleaned_code
    return file_blocks

def interface_fn(prompt):
    try:
        zip_path = generate_code_zip(prompt)
        return zip_path
    except Exception as e:
        return f"Error: {e}"

with gr.Blocks() as demo:
    gr.Markdown("## AI Code Packager (GGUF-Powered)")
    prompt = gr.Textbox(label="Enter AI response or code instructions")
    output = gr.File(label="Download Zip")
    btn = gr.Button("Generate Zip")
    btn.click(fn=interface_fn, inputs=prompt, outputs=output)

demo.launch()