afarovruslan/generate-comments

1import streamlit as st2from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, TextIteratorStreamer3from threading import Thread4import torch5 6# Загрузка модели и токенизатора (кешируем в Streamlit)7@st.cache_resource8def load_model():9    model_name = "afarovruslan/generate-comments-model"10    config = AutoConfig.from_pretrained(model_name)11    tokenizer = AutoTokenizer.from_pretrained(model_name)12    model = AutoModelForCausalLM.from_pretrained(model_name, config=config)13    return model, tokenizer14 15model, tokenizer = load_model()16 17# Интерфейс Streamlit18st.title("GPT-2 Large в реальном времени")19user_input = st.text_input("Введите запрос:", "Расскажи про искусственный интеллект")20 21if st.button("Сгенерировать"):22    # Подготовка потока вывода23    st.markdown("**Результат:**")24    text_output = st.empty()  # Контейнер для динамического обновления25    26    # Инициализация streamer27    streamer = TextIteratorStreamer(28        tokenizer, 29        skip_prompt=True,  # Пропускаем повтор входного запроса30        skip_special_tokens=True31    )32    33    # Подготовка входных данных34    inputs = tokenizer([user_input], return_tensors="pt").to(model.device)35    36    # Запуск генерации в отдельном потоке37    generation_kwargs = dict(38        inputs,39        streamer=streamer,40        max_new_tokens=500,41        do_sample=True,42        top_p=0.95,43        temperature=0.744    )45    thread = Thread(target=model.generate, kwargs=generation_kwargs)46    thread.start()47    48    # Постепенный вывод сгенерированного текста49    full_response = ""50    for new_token in streamer:51        full_response += new_token52        text_output.markdown(full_response + "▌")  # Курсор для эффекта печати53    54    text_output.markdown(full_response)  # Финальный текст без курсора