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

print("🚀 Loading DeepSeek-Coder-1.3B model...")
model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.float32
)
print("✅ Model loaded successfully!")

def generate_code(prompt, max_length=512, temperature=0.7):
    # Fixed: Using proper string concatenation instead of f-string
    instruction_part = "### Instruction:"
    response_part = ""
    formatted_prompt = instruction_part + prompt + response_part
    
    inputs = tokenizer(formatted_prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_length,
            temperature=temperature,
            do_sample=True,
            top_p=0.95,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Extract only the generated response
    if "### Response:" in response:
        return response.split("### Response:")[-1].strip()
    return response

demo = gr.Interface(
    fn=generate_code,
    inputs=[
        gr.Textbox(
            label="Coding Prompt", 
            placeholder="Enter your coding question or request...",
            lines=5
        ),
        gr.Slider(
            minimum=128, 
            maximum=1024, 
            value=512, 
            step=64, 
            label="Max Length"
        ),
        gr.Slider(
            minimum=0.1, 
            maximum=2.0, 
            value=0.7, 
            step=0.1, 
            label="Temperature"
        )
    ],
    outputs=gr.Textbox(label="Generated Code", lines=15),
    title="🚀 DeepSeek Coder 1.3B - Docker Edition",
    description="AI coding assistant powered by DeepSeek-Coder-1.3B running in Docker",
    examples=[
        ["Write a Python function for binary search"],
        ["Create a Flask REST API for user login"],
        ["Write a JavaScript function to debounce API calls"],
        ["Create a SQL query to find top 10 users by activity"]
    ],
    theme=gr.themes.Soft()
)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860,
        show_error=True
    )