# 1. Импорт необходимых библиотек
import gradio as gr
from transformers import pipeline

# 2. Загрузка двух нейросетей
generator = pipeline("text2text-generation", model="google/flan-t5-small")
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ru-en")

# 3. Создание функции для генерации ответа
def generate_response(prompt, language, temperature, max_length, history):
    source_lang_code = "ru" if language == "ru" else "en"
    target_lang_code = "ru" if language == "ru" else "en"
    
    english_prompt = translator(prompt, src_lang=source_lang_code, tgt_lang="en")[0]['translation_text']
    
    response_en = generator(
        english_prompt, 
        max_length=max_length, 
        num_return_sequences=1,
        do_sample=True,  
        temperature=temperature, 
        no_repeat_ngram_size=2  
    )[0]['generated_text']

    final_response = translator(response_en, src_lang="en", tgt_lang=target_lang_code)[0]['translation_text']

    history.append([prompt, final_response])
    
    return history

# 4. Создание веб-интерфейса
with gr.Blocks(theme="soft", title="COLIN") as iface:
    gr.Markdown("<h1>COLIN</h1>")
    gr.Markdown("<h3>Ask questions in English or Russian.</h3>")
    
    with gr.Row():
        language_dropdown = gr.Dropdown(
            ["en", "ru"],
            value="en",
            label="Select Language"
        )
    
    chatbot = gr.Chatbot(label="Chat") # Исправлено здесь
    textbox = gr.Textbox(placeholder="Enter your message here...")
    
    with gr.Row():
        gr.Markdown("<h3>Настройки ответа</h3>")
        temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.7, label="Temperature (Креативность)")
        max_length_slider = gr.Slider(minimum=10, maximum=100, step=10, value=50, label="Max Length (Длина ответа)")

    with gr.Row():
        clear_btn = gr.Button("New Chat")
        submit_btn = gr.Button("Submit")

    submit_btn.click(
        fn=generate_response,
        inputs=[textbox, language_dropdown, temperature_slider, max_length_slider, chatbot],
        outputs=[chatbot]
    )
    clear_btn.click(
        fn=lambda: [],
        inputs=[],
        outputs=[chatbot]
    )

iface.launch()