app.py

import os
from threading import Thread
from typing import Iterator

import gradio as gr
import spaces
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.
"""

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()

@spaces.GPU(duration=90)
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()

    outputs = []
    for text in streamer:
        outputs.append(text)
        yield message + "".join(outputs)

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):
            input_box = gr.Textbox(
                label="Enter your text",
                placeholder="Type your message here...",
                lines=5
            )
            output_box = gr.Textbox(
                label="Completed text",
                lines=10,
                interactive=False
            )
        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,
            )

    complete_btn = gr.Button("Complete")

    complete_btn.click(
        fn=generate,
        inputs=[
            input_box,
            max_new_tokens,
            temperature,
            top_p,
            top_k,
            repetition_penalty
        ],
        outputs=output_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=input_box
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch()