import gradio as gr
import base64
import vertexai
from vertexai.generative_models import GenerativeModel, Part, FinishReason
import vertexai.preview.generative_models as generative_models

# Инициализация Vertex AI
vertexai.init(project="carbide-network-426214-b9", location="us-central1")

# Создание модели
model = GenerativeModel(
    "gemini-1.5-pro-001",
)

# Конфигурации генерации и настройки безопасности
generation_config = {
    "max_output_tokens": 8192,
    "temperature": 1,
    "top_p": 0.95,
}

safety_settings = {
    generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
    generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
    generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
    generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
}

# Функция генерации текста
def generate(prompt):
    responses = model.generate_content(
        [prompt],
        generation_config=generation_config,
        safety_settings=safety_settings,
        stream=True,
    )
    output_text = "".join(response.text for response in responses)
    return output_text

# Создание интерфейса Gradio
with gr.Blocks() as demo:
    gr.Markdown("### Интерфейс для Hugging Face с использованием Gradio и Vertex AI")
    prompt = gr.Textbox(label="Введите ваш запрос", placeholder="Привет! Кто ты?")
    output = gr.Textbox(label="Ответ модели")
    generate_button = gr.Button("Сгенерировать")

    generate_button.click(fn=generate, inputs=prompt, outputs=output)

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