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: | |
| print("\n--- Final Raw AI Response ---") | |
| print(full_raw_response) | |
| print("--------------------------\n") | |
| 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 for a coffee shop with sections for menu, about, and contact.", | |
| 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: e.g., 'Use Tailwind CSS for styling', 'Make it responsive'", | |
| 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="plaintext", | |
| 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) |