Spaces:
Running
Running
File size: 2,270 Bytes
0358312 361b027 d98e70a 0358312 d98e70a 74f7421 d98e70a f0a222c d98e70a 0358312 d98e70a 74f7421 d98e70a 400b4ef 361b027 c946d47 d98e70a 26c7983 0358312 d98e70a 26c7983 d98e70a 361b027 d98e70a |
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 |
import gradio as gr
import torch
import time
from transformers import AutoTokenizer, AutoModelForCausalLM
# ======================================
# Device setup
# ======================================
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ======================================
# Model config
# ======================================
MODEL_NAME = "Salesforce/codegen-350M-mono"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
# Required for GPT-style models
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.eos_token_id
model.to(DEVICE)
model.eval()
# ======================================
# Generation function
# ======================================
def generate_code(prompt: str) -> str:
if not prompt.strip():
return ""
# Easter egg
if "Oz Labs Were Here" in prompt:
time.sleep(10)
return "We're not in Kansas anymore"
full_prompt = (
"# Generate a complete HTML/CSS/JS website.\n"
"# Return ONLY valid code.\n\n"
f"{prompt}\n"
)
inputs = tokenizer(
full_prompt,
return_tensors="pt",
padding=True,
truncation=True
)
# 🔑 Move inputs to the same device as the model
inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
with torch.no_grad():
output_ids = model.generate(
**inputs,
max_length=1024,
temperature=0.7,
do_sample=True,
top_p=0.95
)
decoded = tokenizer.decode(output_ids[0], skip_special_tokens=True)
# Extract HTML if possible
html_start = decoded.find("<")
return decoded[html_start:] if html_start != -1 else decoded
# ======================================
# Gradio UI
# ======================================
app = gr.Interface(
fn=generate_code,
inputs=gr.Textbox(
lines=6,
placeholder="Describe the website you want to generate...",
label="Website description"
),
outputs=gr.Code(label="Generated HTML / CSS / JS"),
title="Oz AI Website Generator",
description="Describe a website idea. The model returns only HTML/CSS/JS."
)
app.launch()
|