import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration

# Загрузка модели и токенизатора
model_name = "google/flan-t5-large"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

def generate_description(name, location, tags, additional_info):
    input_text = (f"Name: {name} Location: {location} Tags: {tags} "
                  f"Additional Information: {additional_info}")
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True)
    outputs = model.generate(inputs['input_ids'],
      max_length=500,          # Максимальная длина сгенерированного текста
      temperature=0.8,         # Контроль случайности (0.8 для более креативного текста)
      top_k=70,                # Рассматриваем только 70 самых вероятных слов
      top_p=0.9,               # Кумулятивная вероятность 90% (nucleus sampling)
      num_beams=1,             # Отключение beam search (или установите 1 для одного луча)
      do_sample=True)          # Остановка генерации, когда результат удовлетворительный
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Создание интерфейса с помощью Gradio
iface = gr.Interface(
    fn=generate_description,
    inputs=["text", "text", "text", "text"],
    outputs="text",
    title="Generate Description",
    description="Enter name, location, tags, and additional information to generate a description.",
    theme="blank"
)

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