import spaces # Импорт spaces обязательно должен быть до импорта torch
import gradio as gr
import os
import torch
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TextIteratorStreamer,
BitsAndBytesConfig
)
from threading import Thread
from torch.nn.attention import SDPBackend, sdpa_kernel
Проверка наличия CUDA и вывод информации для отладки
if torch.cuda.is_available():
print("CUDA доступна: ", torch.cuda.get_device_name(0))
else:
print(“CUDA не доступна, используем CPU.”)
HF_TOKEN = os.getenv(“HF_TOKEN”, None)
REPO_ID = “nicoboss/DeepSeek-R1-Distill-Qwen-32B-Uncensored”
DESCRIPTION = f’‘’
{REPO_ID}
PLACEHOLDER = f"“”
{REPO_ID}
Задайте свой вопрос...
css = “”"
h1 {
text-align: center;
display: block;
}
“”"
Загрузка токенизатора
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
Загрузка модели с 4-битовой квантовкой для оптимизации использования VRAM
if torch.cuda.is_available():
nf4_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type=“nf4”,
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16
)
model = AutoModelForCausalLM.from_pretrained(
REPO_ID,
quantization_config=nf4_config,
device_map=“auto”
)
else:
model = AutoModelForCausalLM.from_pretrained(REPO_ID, torch_dtype=torch.float32)
Инициализация стримера для генерации текста
streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
@spaces.GPU(duration=130)
def chat(message: str,
history: list = None,
temperature: float = 0.7,
max_new_tokens: int = 512,
top_p: float = 0.95,
top_k: int = 40,
repetition_penalty: float = 1.1,
sys_prompt: str = “”,
progress=gr.Progress(track_tqdm=True)):
try:
if history is None:
history =
messages =
response =
# Формирование списка сообщений
messages.append({“role”: “system”, “content”: sys_prompt})
messages.append({“role”: “user”, “content”: message})
# Подготовка входных данных и перенос тензоров на устройство модели
inputs = tokenizer.apply_chat_template(
history + messages,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt"
)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
generate_kwargs = dict(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
streamer=streamer,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=temperature,
top_k=top_k,
top_p=top_p,
repetition_penalty=repetition_penalty,
pad_token_id=tokenizer.eos_token_id,
)
if temperature == 0:
generate_kwargs['do_sample'] = False
response.append({"role": "assistant", "content": ""})
# Использование ускоренного механизма FLASH_ATTENTION
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
thread = Thread(target=model.generate, kwargs=generate_kwargs)
thread.start()
for text in streamer:
response[-1]["content"] += text
yield response
except Exception as e:
print("Ошибка в chat:", e)
gr.Warning(f"Ошибка: {e}")
yield response
with gr.Blocks(css=css, title=f"{REPO_ID} Chat") as demo:
gr.Markdown(DESCRIPTION)
gr.ChatInterface(
fn=chat,
type=“messages”,
chatbot=gr.Chatbot(height=450, type=“messages”, placeholder=PLACEHOLDER, label=‘Чат’),
additional_inputs=[
gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label=“Temperature”),
gr.Slider(minimum=128, maximum=4096, step=1, value=512, label=“Max new tokens”),
gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label=“Top-p”),
gr.Slider(minimum=0, maximum=100, value=40, step=1, label=“Top-k”),
gr.Slider(minimum=0.0, maximum=2.0, value=1.1, step=0.1, label=“Repetition penalty”),
gr.Textbox(value=“”, label=“System prompt”),
],
examples=[
[‘Как создать базу для человека на Марсе? Кратко.’],
[‘Объясни теорию относительности, как будто мне 8 лет.’],
[‘Сколько будет 9000 * 9000?’],
[‘Напиши шуточное поздравление с днём рождения для моего друга Алекса.’],
[‘Обоснуй, почему пингвин может стать отличным королём джунглей.’]
],
cache_examples=False, # Отключаем кэширование примеров, чтобы избежать ошибок
save_history=True
)
demo.queue().launch(ssr_mode=False)
if name == “main”:
demo.launch()