Spaces:
Sleeping
Sleeping
File size: 3,244 Bytes
0bb94c2 b85b849 54312c4 b85b849 67ab732 b85b849 54312c4 4b1e0ae b85b849 54312c4 4b1e0ae 54312c4 b85b849 54312c4 4b1e0ae 54312c4 0bb94c2 b85b849 0bb94c2 b85b849 0bb94c2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# Initialize model and tokenizer
model_id = "Tesslate/WEBGEN-4B-Preview"
# Load model and tokenizer once during app initialization
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="auto"
)
def generate_code(prompt):
inputs = tok(prompt, return_tensors="pt").to(model.device)
# Generate with streaming
from transformers import TextIteratorStreamer
from threading import Thread
streamer = TextIteratorStreamer(tok, skip_special_tokens=True)
generation_kwargs = dict(
**inputs,
max_new_tokens=10000,
temperature=0.7,
top_p=0.9,
do_sample=True,
pad_token_id=tok.eos_token_id,
streamer=streamer
)
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
generated_text = ""
for new_text in streamer:
generated_text += new_text
# Extract only the code portion (remove prompt and any non-code text)
if "```html" in generated_text:
code_start = generated_text.find("```html") + 7
code_end = generated_text.find("```", code_start)
if code_end != -1:
clean_code = generated_text[code_start:code_end].strip()
else:
clean_code = generated_text[code_start:].strip()
elif "<html" in generated_text or "<!DOCTYPE" in generated_text:
# Find the start of HTML code
html_start = generated_text.find("<!DOCTYPE")
if html_start == -1:
html_start = generated_text.find("<html")
if html_start != -1:
clean_code = generated_text[html_start:].strip()
else:
clean_code = generated_text
else:
clean_code = generated_text
yield clean_code
thread.join()
with gr.Blocks() as demo:
gr.Markdown("# 🧪 Text-to-Code Generator")
gr.Markdown("Generate HTML code from natural language prompts with WEBGEN-4B Preview model")
with gr.Row():
with gr.Column():
prompt_input = gr.Textbox(
label="Prompt",
value="Make a single-file landing page for 'LatticeDB'. Style: modern, generous whitespace, Tailwind, rounded-xl, soft gradients. Sections: navbar, hero (headline + 2 CTAs), features grid, pricing (3 tiers), FAQ accordion, footer. Constraints: semantic HTML, no external JS.",
lines=5,
max_lines=10
)
generate_button = gr.Button("Generate Code")
code_output = gr.Code(label="Generated HTML", language="html", lines=20, interactive=False)
with gr.Column():
html_output = gr.HTML(label="Preview")
# When button is clicked, generate code and update both outputs
generate_button.click(
fn=generate_code,
inputs=prompt_input,
outputs=code_output
).then(
fn=lambda code: code,
inputs=code_output,
outputs=html_output
)
demo.launch() |