Spaces:
Running
Running
| import gradio as gr | |
| from huggingface_hub import InferenceClient | |
| import re | |
| import os # Good practice to import os if needed, though not strictly used here yet | |
| # --- Hugging Face Token (Optional but Recommended) --- | |
| # It's better to use a token, especially for private models or higher rate limits | |
| # from huggingface_hub import login | |
| # login("YOUR_HUGGINGFACE_TOKEN") # Replace with your actual token or set HF_TOKEN env var | |
| # --- Inference Client --- | |
| # Consider adding error handling for client initialization if needed | |
| try: | |
| client = InferenceClient("HuggingFaceH4/zephyr-7b-beta") | |
| except Exception as e: | |
| print(f"Error initializing InferenceClient: {e}") | |
| # Optionally, raise the exception or handle it gracefully in the UI | |
| # For now, we'll let it proceed and potentially fail later if client is None | |
| client = None | |
| # --- Parsing Function --- | |
| def parse_files(raw_response): | |
| """ | |
| Parses filenames and code blocks from the raw AI output. | |
| Assumes format: | |
| filename1.ext | |
| ```lang # Optional code block marker | |
| code for file1 | |
| ``` # Optional code block marker | |
| filename2.ext | |
| code for file2 | |
| ... | |
| """ | |
| if not raw_response: | |
| return [] | |
| # Improved pattern to handle optional code blocks and leading/trailing whitespace | |
| # It looks for a filename line followed by content until the next filename line or end of string. | |
| pattern = re.compile( | |
| r"^\s*([\w\-.\/\\]+\.\w+)\s*\n" # Filename line (must have an extension) | |
| r"(.*?)" # Capture content (non-greedy) | |
| r"(?=\n\s*[\w\-.\/\\]+\.\w+\s*\n|\Z)", # Lookahead for next filename or end of string | |
| re.DOTALL | re.MULTILINE | |
| ) | |
| files = pattern.findall(raw_response) | |
| # Clean up content: remove potential code block markers and extra whitespace | |
| cleaned_files = [] | |
| for name, content in files: | |
| # Remove common code block markers (``` optionally followed by lang) | |
| 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())) | |
| # Handle case where the AI might just output code without filenames | |
| if not cleaned_files and raw_response.strip(): | |
| # Basic check if it looks like code (e.g., contains common HTML/CSS/JS chars) | |
| if any(c in raw_response for c in ['<','>','{','}',';','(',')']): | |
| # Default to index.html if no files parsed but content exists | |
| print("Warning: No filenames found, defaulting to index.html") | |
| lang = "html" # Guess language, default to 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" | |
| # Determine a default filename based on guessed language | |
| 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 | |
| # --- Code Generation Function --- | |
| def generate_code(prompt, backend, system_message, max_tokens, temperature, top_p): | |
| """Generates code using the InferenceClient.""" | |
| if not client: | |
| # Return an error structure if client failed to initialize | |
| return "Error: Inference Client not available.", [] | |
| 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> | |
| <head><title>My Site</title></head> | |
| <body><h1>Hello</h1></body> | |
| </html> | |
| style.css | |
| body {{ | |
| font-family: sans-serif; | |
| }} | |
| script.js | |
| console.log('Hello World!'); | |
| Ensure the code is complete and functional for each file. NO commentary, NO explanations, NO markdown formatting like backticks (```). | |
| Start generating the files now. | |
| """.strip() | |
| messages = [ | |
| {"role": "system", "content": full_sys_msg + ("\n" + system_message if system_message else "")}, | |
| {"role": "user", "content": prompt} | |
| ] | |
| try: | |
| response = client.chat_completion( | |
| messages=messages, | |
| max_tokens=int(max_tokens), # Ensure max_tokens is int | |
| temperature=temperature, | |
| top_p=top_p, | |
| stream=False # Ensure streaming is off for this logic | |
| ) | |
| raw = response.choices[0].message.content | |
| print("\n--- Raw AI Response ---") | |
| print(raw) | |
| print("----------------------\n") | |
| files = parse_files(raw) | |
| return None, files # Return None for error, and the list of files | |
| except Exception as e: | |
| print(f"Error during AI generation: {e}") | |
| return f"Error during AI generation: {e}", [] # Return error message | |
| # --- Gradio Event Handler --- | |
| def on_generate(prompt, backend, system_message, max_tokens, temperature, top_p): | |
| """Callback function for the generate button.""" | |
| error_msg, files = generate_code(prompt, backend, system_message, max_tokens, temperature, top_p) | |
| if error_msg: | |
| # Display error in a single tab if generation failed | |
| error_tab = gr.TabItem(label="Error", children=[gr.Textbox(value=error_msg, label="Generation Error")]) | |
| return gr.Tabs(tabs=[error_tab]) # Return a Tabs component with the error tab | |
| if not files: | |
| # Display message if no files were parsed | |
| no_files_tab = gr.TabItem(label="Output", children=[gr.Textbox(value="AI did not return recognizable file content. Check raw output in console.", label="Result")]) | |
| return gr.Tabs(tabs=[no_files_tab]) # Return a Tabs component with this message | |
| tabs = [] | |
| for name, content in files: | |
| name = name.strip() | |
| content = content.strip() | |
| if not name or not content: # Skip empty names or content | |
| print(f"Skipping file with empty name or content: Name='{name}'") | |
| continue | |
| # Determine language for syntax highlighting | |
| lang = "text" # Default | |
| if name.endswith(".html") or name.endswith(".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") or name.endswith(".bash"): | |
| lang = "bash" | |
| tab_item = gr.TabItem(label=name, elem_id=f"tab_{name.replace('.', '_')}", children=[ # Ensure unique elem_id | |
| gr.Code(value=content, language=lang, label=name) # Add label to Code block | |
| ]) | |
| tabs.append(tab_item) | |
| # *** The Key Fix *** | |
| # Return a new gr.Tabs component instance containing the generated TabItems | |
| return gr.Tabs(tabs=tabs) | |
| # --- Gradio UI Definition --- | |
| with gr.Blocks() as demo: | |
| gr.Markdown("### Website Generator (Static / Flask / Node.js)") | |
| gr.Markdown("Describe the website you want to create. The AI will generate the necessary files.") | |
| with gr.Row(): | |
| prompt = gr.Textbox(label="Describe your website", placeholder="E.g., a simple portfolio site with a contact form", scale=3) | |
| backend = gr.Dropdown(["Static", "Flask", "Node.js"], value="Static", label="Backend Technology", scale=1) | |
| with gr.Accordion("Advanced Options", open=False): | |
| system_message = gr.Textbox(label="Extra instructions for the AI (System Message)", placeholder="Optional: e.g., 'Use Bootstrap 5', 'Prefer functional components in React'", value="") | |
| max_tokens = gr.Slider(minimum=256, maximum=4096, value=1536, step=64, label="Max Tokens (Length)") # Increased max | |
| 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)") | |
| generate_button = gr.Button("✨ Generate Code ✨", variant="primary") | |
| gr.Markdown("#### Generated Files") | |
| # Define the Tabs component placeholder. It will be replaced by the output of on_generate. | |
| out_tabs = gr.Tabs(elem_id="output_tabs") | |
| # Button click action | |
| generate_button.click( | |
| on_generate, | |
| inputs=[prompt, backend, system_message, max_tokens, temperature, top_p], | |
| outputs=[out_tabs], # Output the new Tabs component to replace the placeholder | |
| show_progress="full" # Show progress during generation | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(debug=True) # Use debug=True for more detailed error messages in console |