Spaces:
Running
Running
import gradio as gr | |
from huggingface_hub import InferenceClient | |
import re | |
try: | |
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta") | |
client.timeout = 120 | |
except Exception as e: | |
print(f"Error initializing InferenceClient: {e}") | |
client = None | |
def parse_files(raw_response): | |
if not raw_response: | |
return [] | |
pattern = re.compile( | |
r"^\s*([\w\-.\/\\]+\.\w+)\s*\n" | |
r"(.*?)" | |
r"(?=\n\s*[\w\-.\/\\]+\.\w+\s*\n|\Z)", | |
re.DOTALL | re.MULTILINE | |
) | |
files = pattern.findall(raw_response) | |
cleaned_files = [] | |
for name, content in files: | |
content_cleaned = re.sub(r"^\s*```[a-zA-Z]*\n?", "", content, flags=re.MULTILINE) | |
content_cleaned = re.sub(r"\n?```\s*$", "", content_cleaned, flags=re.MULTILINE) | |
cleaned_files.append((name.strip(), content_cleaned.strip())) | |
if not cleaned_files and raw_response.strip(): | |
if any(c in raw_response for c in ['<','>','{','}',';','(',')']): | |
lang = "html" | |
if "{" in raw_response and "}" in raw_response and ":" in raw_response: | |
lang = "css" | |
elif "function" in raw_response or "const" in raw_response or "let" in raw_response: | |
lang = "javascript" | |
default_filename = "index.html" | |
if lang == "css": | |
default_filename = "style.css" | |
elif lang == "javascript": | |
default_filename = "script.js" | |
cleaned_files.append((default_filename, raw_response.strip())) | |
return cleaned_files | |
def stream_and_parse_code(prompt, backend, system_message, max_tokens, temperature, top_p): | |
if not client: | |
error_msg = "Error: Inference Client not available. Check API token or model name." | |
yield { | |
live_output: gr.update(value=error_msg), | |
final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Error", children=[gr.Textbox(value=error_msg)])]) | |
} | |
return | |
full_sys_msg = f""" | |
You are a code generation AI. Given a prompt, generate the necessary files for a website using the {backend} backend. | |
Always include an index.html file. | |
Respond ONLY with filenames and the raw code for each file. | |
Each file must start with its filename on a new line. Example: | |
index.html | |
<!DOCTYPE html> | |
<html></html> | |
style.css | |
body {{}} | |
script.js | |
console.log("Hello"); | |
Ensure the code is complete. NO commentary, NO explanations, NO markdown formatting like backticks (```). | |
Start generating the files now. | |
""".strip() | |
if system_message: | |
full_sys_msg += "\n\n" + system_message | |
messages = [ | |
{"role": "system", "content": full_sys_msg}, | |
{"role": "user", "content": prompt} | |
] | |
full_raw_response = "" | |
error_occurred = False | |
error_message = "" | |
yield { | |
live_output: gr.update(value="Generating stream..."), | |
final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Generating...")]) | |
} | |
try: | |
stream = client.chat_completion( | |
messages, | |
max_tokens=int(max_tokens), | |
stream=True, | |
temperature=temperature, | |
top_p=top_p | |
) | |
for chunk in stream: | |
token = chunk.choices[0].delta.content | |
if token: | |
full_raw_response += token | |
yield { | |
live_output: gr.update(value=full_raw_response), | |
} | |
except Exception as e: | |
print(f"Error during AI streaming: {e}") | |
error_message = f"Error during AI generation: {e}\n\nPartial Response (if any):\n{full_raw_response}" | |
error_occurred = True | |
yield { | |
live_output: gr.update(value=error_message), | |
final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Error")]) | |
} | |
if error_occurred: | |
final_tabs_update = gr.Tabs(tabs=[ | |
gr.TabItem(label="Error", children=[gr.Textbox(value=error_message, label="Generation Error", lines=10)]) | |
]) | |
else: | |
files = parse_files(full_raw_response) | |
if not files: | |
no_files_msg = "AI finished, but did not return recognizable file content or the response was empty. See raw output above." | |
final_tabs_update = gr.Tabs(tabs=[ | |
gr.TabItem(label="Output", children=[gr.Textbox(value=no_files_msg, label="Result")]) | |
]) | |
yield { live_output: gr.update(value=full_raw_response + "\n\n" + no_files_msg), final_tabs: final_tabs_update } | |
return | |
tabs_content = [] | |
for name, content in files: | |
name = name.strip() | |
content = content.strip() | |
if not name or not content: | |
continue | |
lang = "plaintext" | |
if name.endswith((".html", ".htm")): | |
lang = "html" | |
elif name.endswith(".css"): | |
lang = "css" | |
elif name.endswith(".js"): | |
lang = "javascript" | |
elif name.endswith(".py"): | |
lang = "python" | |
elif name.endswith(".json"): | |
lang = "json" | |
elif name.endswith(".md"): | |
lang = "markdown" | |
elif name.endswith((".sh", ".bash")): | |
lang = "bash" | |
elif name.endswith((".xml", ".xaml", ".svg")): | |
lang = "xml" | |
elif name.endswith(".yaml") or name.endswith(".yml"): | |
lang = "yaml" | |
elem_id = f"tab_{re.sub(r'[^a-zA-Z0-9_-]', '_', name)}" | |
tab_item = gr.TabItem(label=name, elem_id=elem_id, children=[ | |
gr.Code(value=content, language=lang, label=name, interactive=False) | |
]) | |
tabs_content.append(tab_item) | |
if not tabs_content: | |
final_tabs_update = gr.Tabs(tabs=[gr.TabItem(label="Output", children=[gr.Textbox(value="No valid files generated after filtering.", label="Result")])]) | |
else: | |
final_tabs_update = gr.Tabs(tabs=tabs_content) | |
yield { | |
live_output: gr.update(value=full_raw_response if not error_occurred else error_message), | |
final_tabs: final_tabs_update | |
} | |
with gr.Blocks(css=".gradio-container { max-width: 95% !important; }") as demo: | |
gr.Markdown("## WebGen AI β One Prompt β Full Website Generator") | |
gr.Markdown("Generates website code based on your description. Raw output streams live, final files appear in tabs below.") | |
with gr.Row(): | |
with gr.Column(scale=2): | |
prompt = gr.Textbox(label="Describe your website", placeholder="E.g., a simple landing page...", lines=3) | |
backend = gr.Dropdown(["Static", "Flask", "Node.js"], value="Static", label="Backend Technology") | |
with gr.Accordion("Advanced Options", open=False): | |
system_message = gr.Textbox(label="Extra instructions for the AI (System Message)", placeholder="Optional", value="", lines=2) | |
max_tokens = gr.Slider(minimum=256, maximum=4096, value=2048, step=64, label="Max Tokens (Output Length)") | |
temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Temperature (Creativity)") | |
top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (Sampling Focus)") | |
generate_button = gr.Button("β¨ Generate Code β¨", variant="primary") | |
with gr.Column(scale=3): | |
gr.Markdown("#### Live Raw Output Stream") | |
live_output = gr.Code(label="Raw AI Stream", language="text", lines=20, interactive=False) | |
gr.Markdown("---") | |
gr.Markdown("#### Final Generated Files (Tabs)") | |
final_tabs = gr.Tabs(elem_id="output_tabs") | |
generate_button.click(stream_and_parse_code, inputs=[prompt, backend, system_message, max_tokens, temperature, top_p], outputs=[live_output, final_tabs], show_progress="hidden") | |
if __name__ == "__main__": | |
demo.launch(debug=True) |