import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Загрузка модели и токенизатора с учетом FP8 и пользовательского кода
model_name = "deepseek-ai/DeepSeek-V3"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.float8_e4m3fn,  # Используем FP8 для совместимости
    device_map="auto"  # Автоматическое распределение по доступным GPU
)

# Системное сообщение
system_message = (
    "Вы — NelzGPT-A1, дружелюбный и умный помощник, созданный Nelz Studio. "
    "Отвечайте на вопросы вежливо и информативно. "
    "Если вы не знаете ответа, скажите, что не можете помочь."
)

def chat_with_model(message, history):
    # Добавляем системное сообщение в начало диалога
    full_prompt = f"{system_message}\n\nПользователь: {message}\nАссистент:"
    
    # Подготовка входных данных для модели
    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
    
    # Генерация ответа с использованием Multi-Token Prediction (MTP)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=200,  # Максимальная длина ответа
            num_return_sequences=1,  # Один ответ
            temperature=0.7,  # Контроль креативности
            top_p=0.9,  # Контроль разнообразия
            do_sample=True  # Включение сэмплирования
        )
    
    # Декодирование ответа
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Убираем системное сообщение и промпт из ответа
    response = response.replace(full_prompt, "").strip()
    
    return response

# Создание Gradio интерфейса
iface = gr.ChatInterface(
    fn=chat_with_model,
    title="NelzGPT-A1 Chatbot",
    description="Чат с нейросетью NelzGPT-A1, созданной Nelz Studio. Системное сообщение задает контекст диалога."
)

# Запуск интерфейса
iface.launch()