app.py

import gradio as gr
from huggingface_hub import InferenceClient

# Укажи свою модель автодополнения кода (можно заменить на другую)
client = InferenceClient(
    model="deepseek-ai/deepseek-coder-1.3b-base",  # или starcoder2, codegemma, etc.
    token=HF_TOKKEN  # или укажи свой токен явно: token="hf_xxxxxxxxxxxxxxxxx"
)

# Функция генерации кода с потоковым выводом
def complete_code(prompt, max_tokens, temperature, top_p):
    response = ""
    for token in client.text_generation(
        prompt,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=True,
    ):
        response += token  # token — это уже str
        yield 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()