import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Model: DeepSeek Coder
model_name = "deepseek-ai/deepseek-coder-1.3b-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# System prompt for Emalawi19
SYSTEM_PROMPT = """
You are Emalawi19 AI WebCode Assistant.
Your only task is to generate website code.
Output HTML, CSS, and JavaScript only.
Include responsive images and animations where needed.
Do not explain anything.
Do not repeat instructions.
Create modern, semantic, clean code.
"""

def generate_code(user_prompt):
    prompt = SYSTEM_PROMPT + "\nUser: " + user_prompt + "\nAI:"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
    outputs = model.generate(**inputs, max_new_tokens=600, do_sample=True, temperature=0.7)
    code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Extract only code after AI:
    if "AI:" in code:
        code = code.split("AI:")[1].strip()
    return code

iface = gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(lines=4, placeholder="Enter website instructions..."),
    outputs=gr.Textbox(lines=25),
    title="Emalawi19 AI WebCode Generator",
    description="Type instructions to generate HTML, CSS, and JavaScript code for websites."
)

iface.launch()