Spaces:
Sleeping
Sleeping
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()
|