File size: 3,188 Bytes
7bd6e67 a7cccd7 7bd6e67 9f304f1 76a6d8e 79a49f5 205ebd7 4cbd18d e4cd356 4cbd18d e4cd356 205ebd7 4cbd18d e4cd356 205ebd7 67fdcc7 205ebd7 7bd6e67 6e4c8e6 7bd6e67 e8081aa 7bd6e67 e4cd356 76a6d8e 7bd6e67 e4cd356 7bd6e67 e4cd356 7bd6e67 e4cd356 7bd6e67 6e4c8e6 df914b7 7bd6e67 9f304f1 06dfd60 7bd6e67 df914b7 |
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 |
import os
from threading import Thread
from typing import Iterator
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
ACCESS_TOKEN = os.getenv("HF_TOKEN", "")
model_id = "HuggingFaceTB/SmolLM-1.7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="auto",
trust_remote_code=True,
token=ACCESS_TOKEN)
tokenizer = AutoTokenizer.from_pretrained(
model_id,
trust_remote_code=True,
token=ACCESS_TOKEN)
tokenizer.use_default_system_prompt = False
model.config.gradient_checkpointing = True
@spaces.GPU
def generate(
message: str,
system_prompt: str,
max_new_tokens: int = 1024,
temperature: float = 0.01,
top_p: float = 1.00,
) -> Iterator[str]:
conversation = []
if system_prompt:
conversation.append({"role": "system", "content": system_prompt})
conversation.append({"role": "user", "content": message})
input_ids = tokenizer.apply_chat_template(conversation, 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 from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
input_ids = input_ids.to(model.device)
'''
terminators = [
tokenizer.eos_token_id,
tokenizer.convert_tokens_to_ids("<|eot_id|>")
]
'''
streamer = TextIteratorStreamer(tokenizer, timeout=600.0, skip_prompt=True, skip_special_tokens=True)
generate_kwargs = dict(
{"input_ids": input_ids},
streamer=streamer,
max_new_tokens=max_new_tokens,
#eos_token_id=terminators,
do_sample=True,
top_p=top_p,
temperature=temperature,
num_beams=1,
pad_token_id=tokenizer.eos_token_id,
)
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
outputs = []
for text in streamer:
outputs.append(text)
yield "".join(outputs)
chat_interface = gr.Interface(
fn=generate,
inputs=[
gr.Textbox(lines=2, placeholder="Prompt", label="Prompt"),
],
outputs="text",
additional_inputs=[
gr.Textbox(label="System prompt", lines=6),
gr.Slider(
label="Max new tokens",
minimum=1,
maximum=MAX_MAX_NEW_TOKENS,
step=1,
value=DEFAULT_MAX_NEW_TOKENS,
),
gr.Slider(
label="Temperature",
minimum=0.1,
maximum=4.0,
step=0.01,
value=0.01,
),
gr.Slider(
label="Top-p (nucleus sampling)",
minimum=0.05,
maximum=1.0,
step=0.01,
value=1.0,
),
],
title="Model testing - HuggingFaceTB/SmolLM-1.7B-Instruct",
description="Provide system settings and a prompt to interact with the model.",
)
chat_interface.queue(max_size=20).launch()
|