app.py

import gradio as gr
import os
from huggingface_hub import InferenceClient

# Получаем токен из переменной окружения
HF_TOKEN = os.getenv("HF_TOKEN")

# Клиент Hugging Face для выбранной модели
client = InferenceClient(
    model="deepseek-ai/deepseek-coder-1.3b-base",  # Можешь заменить на starcoder2 или codegemma
    token=HF_TOKEN
)

# Функция генерации без stream (deepseek не поддерживает stream=True)
def complete_code(prompt, max_tokens, temperature, top_p):
    response = client.text_generation(
        prompt,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=False  # ⚠️ обязательно False для этой модели
    )
    return response

# Интерфейс Gradio
demo = gr.Interface(
    fn=complete_code,
    inputs=[
        gr.Textbox(lines=10, label="Введите код или комментарий"),
        gr.Slider(minimum=1, maximum=1024, value=128, step=1, label="Максимум новых токенов"),
        gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"),
    ],
    outputs=gr.Textbox(lines=15, label="Дополненный код"),
    title="🛠️ Уязвимый агент автодополнения кода",
    description=(
        "Этот агент демонстрирует поведение LLM без фильтрации: он уязвим к prompt injection, "
        "может генерировать вредоносные команды и не имеет встроенных ограничений. "
        "Использовать только в исследовательских целях!"
    ),
)

if __name__ == "__main__":
    demo.launch()