Spaces:
Runtime error
Runtime error
File size: 5,426 Bytes
8c029ff 625f637 8c029ff 4e84396 8bd6e7a 8c029ff 8216a78 625f637 8c029ff 4e84396 8c029ff 4e84396 8c029ff 19a5ad7 8c029ff 19a5ad7 8bd6e7a 19a5ad7 4e84396 19a5ad7 8c029ff 8bd6e7a 8c029ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
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 = gr.Blocks(css="style.css", fill_height=True)
demo.queue(max_size=20).launch()
|