| import spaces |
| import gradio as gr |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| |
| MODEL_ID = "DeepHat/DeepHat-V1-7B" |
| SYSTEM_PROMPT = ( |
| "You are DeepHat, a helpful assistant specializing in cybersecurity, " |
| "programming, and DevOps. Answer in the user's language. " |
| "Be precise, explain assumptions, and acknowledge uncertainty." |
| ) |
| MAX_INPUT_TOKENS = 4096 |
|
|
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
| |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL_ID, |
| torch_dtype=torch.bfloat16, |
| device_map="cuda", |
| low_cpu_mem_usage=True, |
| attn_implementation="sdpa", |
| ).eval() |
|
|
|
|
| def text_content(content): |
| if isinstance(content, str): |
| return content |
| if isinstance(content, list): |
| return "\n".join( |
| item.get("text", "") |
| for item in content |
| if isinstance(item, dict) and item.get("type") == "text" |
| ) |
| return "" |
|
|
|
|
| def prepare_inputs(message, history): |
| if not message.strip(): |
| raise gr.Error("Введите сообщение.") |
| recent = [] |
| for item in (history or [])[-12:]: |
| if item.get("role") in ("user", "assistant"): |
| content = text_content(item.get("content", "")) |
| if content: |
| recent.append({"role": item["role"], "content": content}) |
| while recent and recent[0]["role"] != "user": |
| recent.pop(0) |
| while True: |
| messages = ( |
| [{"role": "system", "content": SYSTEM_PROMPT}] |
| + recent |
| + [{"role": "user", "content": message}] |
| ) |
| prompt = tokenizer.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True |
| ) |
| inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False) |
| if inputs.input_ids.shape[1] <= MAX_INPUT_TOKENS: |
| return inputs |
| if not recent: |
| raise gr.Error("Сообщение слишком длинное. Сократите его и попробуйте ещё раз.") |
| recent.pop(0) |
| while recent and recent[0]["role"] != "user": |
| recent.pop(0) |
|
|
|
|
| @spaces.GPU(duration=60) |
| def generate(inputs, max_tokens): |
| inputs = inputs.to("cuda") |
| with torch.inference_mode(): |
| output = model.generate( |
| **inputs, |
| max_new_tokens=int(max_tokens), |
| do_sample=True, |
| temperature=0.6, |
| top_p=0.9, |
| repetition_penalty=1.05, |
| pad_token_id=tokenizer.eos_token_id, |
| ) |
| return tokenizer.decode( |
| output[0, inputs.input_ids.shape[1]:], skip_special_tokens=True |
| ).strip() |
|
|
|
|
| def respond(message, history, max_tokens): |
| return generate(prepare_inputs(message, history), max_tokens) |
|
|
|
|
| demo = gr.ChatInterface( |
| fn=respond, |
| type="messages", |
| title="WhiteRabbitNeo / DeepHat 7B", |
| description=( |
| "Помощник по кибербезопасности, программированию и DevOps. " |
| "Ответ появится после завершения генерации. " |
| "ZeroGPU работает с очередью и дневной квотой. " |
| "Длинная история сокращается до последних сообщений." |
| ), |
| additional_inputs=[ |
| gr.Slider(64, 512, value=256, step=64, label="Максимальная длина ответа в токенах") |
| ], |
| examples=[ |
| ["Объясни разницу между XSS и CSRF и способы защиты.", 256], |
| ["Помоги составить план проверки безопасности моего веб-приложения.", 256], |
| ], |
| cache_examples=False, |
| concurrency_limit=1, |
| ) |
|
|
| if __name__ == "__main__": |
| demo.queue(max_size=8).launch() |
|
|
