import os
from threading import Thread
from typing import Iterator

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
DESCRIPTION = """\
# Qwen 0.5B Text Completion

This is a demo of [`Qwen/Qwen2-0.5B-Instruct`](https://huggingface.co/Qwen/Qwen2-0.5B-Instruct), a lightweight language model fine-tuned for instruction following.

This space allows you to input text and have the AI complete it. Simply type your text in the input box, click "Complete", and watch as the AI generates a continuation of your text.

You can adjust various parameters such as temperature and top-p sampling to control the generation process.

Note: You may see a warning about bitsandbytes being compiled without GPU support. This is expected in environments without GPU and does not affect the basic functionality of the demo.
"""

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model_id = "Qwen/Qwen2-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
model.eval()
def generate(
    message: str,
    max_new_tokens: int = 1024,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> Iterator[str]:
    input_ids = tokenizer.encode(message, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        {"input_ids": input_ids},
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        num_beams=1,
        repetition_penalty=repetition_penalty,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    full_message = message
    for text in streamer:
        full_message += text
        yield full_message

with gr.Blocks(css="style.css", fill_height=True) as demo:
            gr.Markdown(DESCRIPTION)
            gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")

            with gr.Row():
                with gr.Column(scale=4):
                    text_box = gr.Textbox(
                        label="Enter your text",
                        placeholder="Type your message here...",
                        lines=10
                    )
                with gr.Column(scale=1):
                    max_new_tokens = gr.Slider(
                        label="Max new tokens",
                        minimum=1,
                        maximum=MAX_MAX_NEW_TOKENS,
                        step=1,
                        value=DEFAULT_MAX_NEW_TOKENS,
                    )
                    temperature = gr.Slider(
                        label="Temperature",
                        minimum=0.1,
                        maximum=4.0,
                        step=0.1,
                        value=0.6,
                    )
                    top_p = gr.Slider(
                        label="Top-p (nucleus sampling)",
                        minimum=0.05,
                        maximum=1.0,
                        step=0.05,
                        value=0.9,
                    )
                    top_k = gr.Slider(
                        label="Top-k",
                        minimum=1,
                        maximum=1000,
                        step=1,
                        value=50,
                    )
                    repetition_penalty = gr.Slider(
                        label="Repetition penalty",
                        minimum=1.0,
                        maximum=2.0,
                        step=0.05,
                        value=1.2,
                    )

            with gr.Row():
                complete_btn = gr.Button("Complete")
                stop_btn = gr.Button("Stop Generation")

            stop_click = stop_btn.click(fn=None, cancels=[complete_btn.click])

            complete_btn.click(
                fn=generate,
                inputs=[
                    text_box,
                    max_new_tokens,
                    temperature,
                    top_p,
                    top_k,
                    repetition_penalty
                ],
                outputs=text_box
            )

            gr.Examples(
                examples=[
                    "Hello there! How are you doing?",
                    "Can you explain briefly to me what is the Python programming language?",
                    "Explain the plot of Cinderella in a sentence.",
                    "How many hours does it take a man to eat a Helicopter?",
                    "Write a 100-word article on 'Benefits of Open-Source in AI research'",
                ],
                inputs=text_box
            )
            if __name__ == "__main__":
                demo.queue(max_size=20).launch()