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

title = """# 🤵Welcome to SAK's CodBot 💻 (A Coding ChatBot) ⭐⭐⭐"""

MARKDOWN = """
CodingBot is an open-source coding language model that delivers excellent performance.

CodingBot leverages LlamaforCausalLM.

CodingBot supports 'java', 'javascript', 'c++', 'c#', 'c', 'html', 'java_server_pages', 'python', 'php', 'go', 'kotlin', 'swift', 'dart', 'shell', 'json', 'lua', 'matlab', 'yaml', 'css', 'rust', 'sql', 'ruby', 'tex', 'objective-c', 'powershell', 'ocaml', 'groovy', 'cmake', 'julia', 'perl', 'assembly', 'haskell', 'fortran', 'pascal', 'rmarkdown', 'scala', 'visual_basic', 'verilog', 'prolog', 'r', 'dockerfile','cobol', 'batchfile', 'toml', 'lisp', 'erlang', 'coffeescript', 'makefile', 'clojure', 'elixir'

**Demo by [Sunder Ali Khowaja](https://sander-ali.github.io) - [X](https://x.com/SunderAKhowaja) -[Github](https://github.com/sander-ali) -[Hugging Face](https://huggingface.co/SunderAli17)**
"""

# Define the device and model path
device = "cuda" if torch.cuda.is_available() else "cpu"
model_path = "SunderAli17/CodingBot"

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto").eval()

@spaces.GPU(duration=130)
def generate_code(system_prompt, user_prompt, max_length):
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(device)

    generated_ids = model.generate(
        model_inputs.input_ids,
        max_new_tokens=max_length,
        eos_token_id=tokenizer.eos_token_id  
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return response

theme = gr.themes.Soft(
    font=[gr.themes.GoogleFont('Source Code Pro'), gr.themes.GoogleFont('Public Sans'), 'system-ui', 'sans-serif'],
)
js_func = """
function refresh() {
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""
with gr.Blocks(js = js_func, theme = theme) as SAK:
    gr.Markdown(title)
    gr.Markdown(MARKDOWN)

    system_prompt_input = gr.Textbox(
        label="👨‍💻 CodBot Instruction:",
        value="Hello Sir! how are you today? I am here to generate clear and concise code examples for you.",
        lines=2
    )
    user_prompt_input = gr.Code(
        label="❓ Coding Prompt 💻",
        value="Shopping website in HTML and Java",
        language="python",
        lines=2
    )
    code_output = gr.Code(label="👨‍💻 CodBot", language='python', lines=50, interactive=True)
    max_length_slider = gr.Slider(minimum=1, maximum=1800, value=650, label="Max Token Length")

    generate_button = gr.Button("Generate Code")
    generate_button.click(
        generate_code,
        inputs=[system_prompt_input, user_prompt_input, max_length_slider],
        outputs=code_output
    )

SAK.queue().launch(debug=True, share=True)