import torch,gradio as gr,spaces
from transformers import AutoTokenizer,AutoModelForCausalLM
MODEL_ID="caikybaldo999/webcoder-100m-html-css-js"
tok=AutoTokenizer.from_pretrained(MODEL_ID)
model=None
@spaces.GPU(duration=60)
def generate(prompt):
    global model
    if model is None: model=AutoModelForCausalLM.from_pretrained(MODEL_ID,torch_dtype=torch.bfloat16).to("cuda").eval()
    p=f"<|user|>\n{prompt}<|end|>\n<|assistant|>\n"
    x=tok(p,return_tensors="pt").to("cuda")
    with torch.inference_mode(): y=model.generate(**x,max_new_tokens=768,do_sample=True,temperature=.7,top_p=.92,pad_token_id=tok.pad_token_id)
    return tok.decode(y[0,x["input_ids"].shape[1]:],skip_special_tokens=True)
demo=gr.Interface(fn=generate,inputs=gr.Textbox(lines=8,label="Pedido"),outputs=gr.Code(label="Código"),title="WebCoder 100M",description="HTML • CSS • JavaScript")
demo.queue().launch()