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

# Настройки
MODEL_NAME = "deepseek-ai/deepseek-llm-7b-chat"

# Проверяем доступность CUDA
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Используемое устройство: {DEVICE}")

# Загружаем токенизатор и модель
print("Загрузка токенизатора...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

print("Загрузка модели...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Используем float16 для GPU
    device_map="auto"  # Автоматическое распределение на GPU
)

# Устанавливаем pad_token (иначе может зависать)
model.config.pad_token_id = tokenizer.eos_token_id

# Функция генерации текста
def chat_with_deepseek(prompt, max_length=512, temperature=0.7):
    if not prompt.strip():
        return "Введите текст запроса!"

    print(f"Запрос: {prompt}")
    
    # Токенизируем ввод и перемещаем на устройство
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            temperature=temperature
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    print(f"Ответ модели: {response}")
    return response

# Gradio UI
iface = gr.Interface(
    fn=chat_with_deepseek,
    inputs=[
        gr.Textbox(label="Введите запрос"),
        gr.Slider(50, 1024, step=50, label="Макс. длина"),
        gr.Slider(0.1, 1.5, step=0.1, label="Температура")
    ],
    outputs=gr.Textbox(label="Ответ DeepSeek"),
    title="🤖 DeepSeek AI Chat",
    description="Отвечает на любые вопросы с помощью DeepSeek LLM."
)

# Запуск
if __name__ == "__main__":
    iface.launch()