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

# Model
model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# Load model (CPU optimized)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True
)

model.to("cpu")
model.eval()

# Generate function
def generate_code(prompt):
    if not prompt.strip():
        return "Please enter a prompt."

    formatted_prompt = f"""You are a professional programmer.
Write clean, complete, and correct code.

Instruction:
{prompt}

Response:
"""

    inputs = tokenizer(
        formatted_prompt,
        return_tensors="pt",
        truncation=True,
        max_length=512
    )

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=120,
            temperature=0.2,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.2,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id
        )

    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Clean output
    result = output_text.replace(formatted_prompt, "").strip()

    return result

# UI
iface = gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(
        lines=5,
        placeholder="Example: Create a login page using HTML and CSS"
    ),
    outputs=gr.Textbox(lines=12),
    title="DeepSeek Coder AI (Optimized)",
    description="Code generator running on Hugging Face Spaces (CPU optimized)."
)

iface.launch()