Spaces:
Running
Running
| import gradio as gr | |
| from huggingface_hub import InferenceClient | |
| import os | |
| import re | |
| API_TOKEN = os.getenv("HF_TOKEN", None) | |
| MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct" | |
| try: | |
| client = InferenceClient(model=MODEL, token=API_TOKEN) if API_TOKEN else InferenceClient(model=MODEL) | |
| except Exception as e: | |
| raise gr.Error(f"Failed to initialize model client. Error: {e}") | |
| def generate_code(prompt: str, backend_choice: str, max_tokens: int, temperature: float, top_p: float): | |
| system_message = ( | |
| "You are an AI assistant programmed to generate website codes only. " | |
| "You must not use triple backticks (```html, ```python, etc.). " | |
| "If multiple files are needed, separate them clearly using:\n" | |
| "TAB.NAME={filename}\n" | |
| "Only generate code. No explanations, no phrases like 'Here is the code'. " | |
| "If user asks non-website code, reply:\n" | |
| "'hey there! am here to create websites for you unfortunately am programmed to not create codes! otherwise I would go on the naughty list :-('." | |
| ) | |
| user_prompt = f"USER_PROMPT = {prompt}\nUSER_BACKEND = {backend_choice}" | |
| messages = [ | |
| {"role": "system", "content": system_message}, | |
| {"role": "user", "content": user_prompt} | |
| ] | |
| full_response = "" | |
| try: | |
| stream = client.chat_completion( | |
| messages=messages, | |
| max_tokens=max_tokens, | |
| stream=True, | |
| temperature=temperature, | |
| top_p=top_p, | |
| ) | |
| for message in stream: | |
| token = message.choices[0].delta.content | |
| if isinstance(token, str): | |
| full_response += token | |
| cleaned_response = full_response.strip() | |
| cleaned_response = re.sub(r"^\s*```[a-z]*\s*\n?", "", cleaned_response) | |
| cleaned_response = re.sub(r"\n?\s*```\s*$", "", cleaned_response) | |
| cleaned_response = re.sub(r"<\s*\|?\s*(user|assistant|system|endoftext)\s*\|?\s*>", "", cleaned_response, flags=re.IGNORECASE) | |
| cleaned_response = cleaned_response.replace("<|im_end|>", "").replace("<|im_start|>", "").strip() | |
| common_phrases = [ | |
| "Here is the code:", "Okay, here is the code:", "Here's the code:", | |
| "Sure, here is the code you requested:", "Let me know if you need anything else.", | |
| "Here is the website code you requested:", "Here are the files for your website:", | |
| "Okay, here are the files:" | |
| ] | |
| lower_response = cleaned_response.lower() | |
| for phrase in common_phrases: | |
| if lower_response.startswith(phrase.lower()): | |
| cleaned_response = cleaned_response[len(phrase):].lstrip() | |
| lower_response = cleaned_response.lower() | |
| if not cleaned_response: | |
| return "Error: Empty response from model after cleaning." | |
| return cleaned_response | |
| except Exception as e: | |
| return f"## Error\n\nFailed to generate code.\n**Reason:** {e}" | |
| def split_files(full_code_text): | |
| file_blocks = [] | |
| splits = re.split(r'(TAB\.NAME=\{.+?\})', full_code_text) | |
| initial_content = splits[0].strip() | |
| if len(splits) == 1: | |
| if initial_content: | |
| default_name = "index.html" | |
| if "def " in initial_content or "import " in initial_content: default_name = "app.py" | |
| elif "function " in initial_content or "const " in initial_content or "let " in initial_content: default_name = "script.js" | |
| elif "<!DOCTYPE html>" in initial_content or "<html" in initial_content: default_name = "index.html" | |
| elif "@app.route" in initial_content: default_name = "app.py" | |
| elif "require(" in initial_content or "module.exports" in initial_content: default_name = "server.js" | |
| elif "<?php" in initial_content: default_name = "index.php" | |
| elif "package main" in initial_content: default_name = "main.go" | |
| file_blocks.append((default_name, initial_content)) | |
| else: | |
| for i in range(1, len(splits), 2): | |
| marker = splits[i] | |
| content = splits[i+1].strip() if (i+1) < len(splits) else "" | |
| filename_match = re.search(r'TAB\.NAME=\{(.+?)\}', marker) | |
| if filename_match: | |
| filename = filename_match.group(1).strip() | |
| if content: | |
| file_blocks.append((filename, content)) | |
| elif i == 1 and initial_content: # Handle content before the first explicit marker | |
| file_blocks.append(("file_0.txt", initial_content)) | |
| if content: # Add the content after the first marker if it exists | |
| filename_match_fallback = re.search(r'TAB\.NAME=\{(.+?)\}', marker) | |
| if filename_match_fallback: | |
| filename = filename_match_fallback.group(1).strip() | |
| file_blocks.append((filename, content)) | |
| return file_blocks | |
| with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo: | |
| gr.Markdown("# ✨ Website Code Generator ✨") | |
| gr.Markdown( | |
| "Describe the website you want. The AI will generate website code.\n\n" | |
| "**Rules:**\n" | |
| "- Provide a backend hint (Static / Flask / Node.js).\n" | |
| "- Generated code should be functional and SFW.\n" | |
| "- Only generates website-related code.\n" | |
| "- If multiple files are generated, they will be separated below using the format `TAB.NAME={filename}`." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| prompt_input = gr.Textbox( | |
| label="Website Description", | |
| placeholder="e.g., A simple Flask app with one route that displays 'Hello World'.", | |
| lines=6, | |
| ) | |
| backend_radio = gr.Radio( | |
| ["Static", "Flask", "Node.js"], | |
| label="Backend Context / Hint", | |
| value="Static", | |
| ) | |
| generate_button = gr.Button("✨ Generate Website Code", variant="primary") | |
| with gr.Column(scale=3): | |
| main_output_label = gr.Markdown("### Full Generated Code / Main File") | |
| main_output_code = gr.Code( | |
| label="Generated Code", # Label is less prominent now | |
| language="html", | |
| lines=15, | |
| interactive=False, | |
| ) | |
| extra_outputs_column = gr.Column(visible=False) # Initially hidden | |
| with gr.Accordion("Advanced Settings", open=False): | |
| max_tokens_slider = gr.Slider( | |
| minimum=512, | |
| maximum=8192, | |
| value=4096, | |
| step=256, | |
| label="Max New Tokens" | |
| ) | |
| temperature_slider = gr.Slider( | |
| minimum=0.1, | |
| maximum=1.2, | |
| value=0.7, | |
| step=0.1, | |
| label="Temperature" | |
| ) | |
| top_p_slider = gr.Slider( | |
| minimum=0.1, | |
| maximum=1.0, | |
| value=0.9, | |
| step=0.05, | |
| label="Top-P" | |
| ) | |
| def get_language(filename): | |
| if filename.endswith(".html") or filename.endswith(".htm"): return "html" | |
| if filename.endswith(".css"): return "css" | |
| if filename.endswith(".js"): return "javascript" | |
| if filename.endswith(".py"): return "python" | |
| if filename.endswith(".json"): return "json" | |
| if filename.endswith(".sql"): return "sql" | |
| if filename.endswith(".php"): return "php" | |
| if filename.endswith(".go"): return "go" | |
| if filename.endswith(".java"): return "java" | |
| if filename.endswith(".rb"): return "ruby" | |
| if filename.endswith(".sh"): return "shell" | |
| if filename.endswith(".yml") or filename.endswith(".yaml"): return "yaml" | |
| if filename.endswith(".md"): return "markdown" | |
| return "text" | |
| def generate_and_display(prompt, backend, max_tokens, temperature, top_p): | |
| full_code = generate_code(prompt, backend, max_tokens, temperature, top_p) | |
| if full_code.startswith("## Error"): | |
| return { | |
| main_output_label: gr.Markdown.update(value="### Error Occurred"), | |
| main_output_code: gr.Code.update(value=full_code, language="markdown"), | |
| extra_outputs_column: gr.Column.update(visible=False, children=[]) | |
| } | |
| if full_code.startswith("Error: Empty response"): | |
| return { | |
| main_output_label: gr.Markdown.update(value="### Error Occurred"), | |
| main_output_code: gr.Code.update(value=full_code, language="text"), | |
| extra_outputs_column: gr.Column.update(visible=False, children=[]) | |
| } | |
| files = split_files(full_code) | |
| dynamic_components = [] | |
| main_file_content = full_code | |
| main_file_lang = "text" | |
| main_file_label = "### Full Generated Code" | |
| if not files: | |
| main_file_content = full_code # Show full code if split failed but we got output | |
| main_file_lang = get_language("output.txt") # Basic guess | |
| main_file_label = "### Full Generated Output (No Files Detected)" | |
| return { | |
| main_output_label: gr.Markdown.update(value=main_file_label), | |
| main_output_code: gr.Code.update(value=main_file_content, language=main_file_lang), | |
| extra_outputs_column: gr.Column.update(visible=False, children=[]) | |
| } | |
| if len(files) == 1: | |
| main_file_content = files[0][1] | |
| main_file_lang = get_language(files[0][0]) | |
| main_file_label = f"### File: {files[0][0]}" | |
| return { | |
| main_output_label: gr.Markdown.update(value=main_file_label), | |
| main_output_code: gr.Code.update(value=main_file_content, language=main_file_lang), | |
| extra_outputs_column: gr.Column.update(visible=False, children=[]) | |
| } | |
| else: | |
| main_file_content = files[0][1] | |
| main_file_lang = get_language(files[0][0]) | |
| main_file_label = f"### File: {files[0][0]}" | |
| for i, (filename, content) in enumerate(files[1:], start=1): # Start from the second file for extras | |
| lang = get_language(filename) | |
| dynamic_components.append(gr.Markdown(f"### File: {filename}")) | |
| dynamic_components.append(gr.Code(value=content, language=lang, label=filename, interactive=False)) | |
| return { | |
| main_output_label: gr.Markdown.update(value=main_file_label), | |
| main_output_code: gr.Code.update(value=main_file_content, language=main_file_lang), | |
| extra_outputs_column: gr.Column.update(visible=True, children=dynamic_components) | |
| } | |
| generate_button.click( | |
| fn=generate_and_display, | |
| inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider], | |
| outputs=[main_output_label, main_output_code, extra_outputs_column] | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch() |