MINEOGO's picture
Update app.py
283c557 verified
raw
history blame
7.87 kB
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)