Update app.py

app.py

@@ -2,128 +2,146 @@ import streamlit as st
 import torch
 import transformers
 from trl import AutoModelForCausalLMWithValueHead
+import math
+import time
 
-# Настройки страницы
-st.set_page_config(page_title="RLHF: IMDB Movie Reviews", layout="wide")
-st.title("🎬 Генерация отзывов на фильмы с помощью RLHF")
+# ============================================================
+# НАСТРОЙКИ СТРАНИЦЫ И СТИЛИ (Вау-эффект)
+# ============================================================
+st.set_page_config(page_title="RLHF Magic | Movie Reviews", page_icon="🍿", layout="wide")
+
+# Делаем кастомный CSS для красоты
+st.markdown("""
+<style>
+    .big-font { font-size:22px !important; font-weight: 500; }
+    .stProgress .st-bo { transition: background-color 0.5s ease; }
+</style>
+""", unsafe_allow_html=True)
+
+st.title("🍿 Нейросеть-Кинокритик: До и После RLHF")
 st.markdown("""
-Это приложение сравнивает два варианта модели:
-- **Original GPT-2**: базовая модель, обученная на отзывах IMDB.
-- **RLHF Model (PPO)**: та же модель, но дообученная с помощью RLHF писать **только позитивные** отзывы.
-""")
+<div class="big-font">
+Посмотрите, как работает магия обучения с подкреплением (RLHF). <br>
+Слева — базовая модель GPT-2, которая пишет что вздумается. Справа — та же модель, но <b>натренированная всегда писать позитивные отзывы</b>, даже если вы начинаете текст с ужасных слов!
+</div>
+<br>
+""", unsafe_allow_html=True)
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # ============================================================
-# ЗАГРУЗКА МОДЕЛЕЙ (кешируем, чтобы не грузить при каждом нажатии)
+# ЗАГРУЗКА МОДЕЛЕЙ
 # ============================================================
 @st.cache_resource
 def load_models():
-    with st.spinner("Загрузка моделей в память... Пожалуйста, подождите (это делается 1 раз)."):
-        # 1. Загрузка Reward Model
-        reward_path = "reward_model_trained"
-        reward_tokenizer = transformers.AutoTokenizer.from_pretrained(reward_path)
-        reward_model = transformers.AutoModelForSequenceClassification.from_pretrained(reward_path).to(DEVICE).eval()
-        
-        # 2. Загрузка Original Model (Базовая до RLHF)
-        orig_model_name = "lvwerra/gpt2-imdb"
-        orig_tokenizer = transformers.AutoTokenizer.from_pretrained(orig_model_name)
-        if orig_tokenizer.pad_token is None:
-            orig_tokenizer.pad_token = orig_tokenizer.eos_token
-            
-        orig_model = transformers.AutoModelForCausalLM.from_pretrained(orig_model_name).to(DEVICE).eval()
-        
-        # 3. Загрузка RLHF Model (Обученная через PPO)
-        ppo_path = "ppo_model_trained"
-        # Для генерации нам нужен только CausalLM, но чтобы загрузить веса корректно, используем ValueHead класс
-        rlhf_model_full = AutoModelForCausalLMWithValueHead.from_pretrained(ppo_path).to(DEVICE).eval()
-        rlhf_model = rlhf_model_full.pretrained_model # вытаскиваем саму языковую модель
-        
-        return reward_model, reward_tokenizer, orig_model, orig_tokenizer, rlhf_model
+    reward_path = "reward_model_trained"
+    ppo_path = "ppo_model_trained"
+    orig_model_name = "lvwerra/gpt2-imdb"
+    
+    # 1. Reward Model
+    reward_tokenizer = transformers.AutoTokenizer.from_pretrained(reward_path)
+    reward_model = transformers.AutoModelForSequenceClassification.from_pretrained(reward_path).to(DEVICE).eval()
+    
+    # 2. Original GPT-2
+    orig_tokenizer = transformers.AutoTokenizer.from_pretrained(orig_model_name)
+    if orig_tokenizer.pad_token is None:
+        orig_tokenizer.pad_token = orig_tokenizer.eos_token
+    orig_model = transformers.AutoModelForCausalLM.from_pretrained(orig_model_name).to(DEVICE).eval()
+    
+    # 3. RLHF Model
+    rlhf_model_full = AutoModelForCausalLMWithValueHead.from_pretrained(ppo_path).to(DEVICE).eval()
+    rlhf_model = rlhf_model_full.pretrained_model 
+    
+    return reward_model, reward_tokenizer, orig_model, orig_tokenizer, rlhf_model
 
-try:
+with st.spinner("⏳ Подготовка нейросетей... (занимает около минуты при первом старте)"):
     reward_model, reward_tokenizer, orig_model, orig_tokenizer, rlhf_model = load_models()
-except Exception as e:
-    st.error(f"Ошибка загрузки моделей! Убедитесь, что па��ки `reward_model_trained` и `ppo_model_trained` находятся рядом с app.py.\nДетали: {e}")
-    st.stop()
 
 # ============================================================
-# ФУНКЦИИ ГЕНЕРАЦИИ И ОЦЕНКИ
+# ФУНКЦИИ МАГИИ
 # ============================================================
 def compute_reward(text):
-    """Вычисляет 'позитивность' текста с помощью Reward модели"""
-    inputs = reward_tokenizer(
-        text, truncation=True, max_length=512,
-        padding=True, return_tensors="pt"
-    ).to(DEVICE)
+    inputs = reward_tokenizer(text, truncation=True, max_length=512, padding=True, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
         score = reward_model(**inputs).logits[0, 0].item()
     return score
 
+# Функция перевода Reward score в проценты (Sigmoid)
+def get_positivity_percent(score):
+    return int((1 / (1 + math.exp(-score))) * 100)
+
 def generate_text(model, tokenizer, prompt, max_new_tokens, temperature, top_p):
-    """Генерирует продолжение текста"""
     inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            do_sample=True,
-            temperature=temperature,
-            top_p=top_p,
-            pad_token_id=tokenizer.eos_token_id
-        )
+        outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=True, 
+                                 temperature=temperature, top_p=top_p, pad_token_id=tokenizer.eos_token_id)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+# Генератор для эффекта печатной машинки
+def stream_text(text, delay=0.03):
+    for word in text.split(" "):
+        yield word + " "
+        time.sleep(delay)
+
 # ============================================================
-# ИНТЕРФЕЙС ПРИЛОЖЕНИЯ
+# ИНТЕРФЕЙС И ВЗАИМОДЕЙСТВИЕ
 # ============================================================
-st.sidebar.header("Параметры генерации")
-max_tokens = st.sidebar.slider("Max New Tokens", 10, 150, 80)
-temperature = st.sidebar.slider("Temperature", 0.1, 1.5, 0.8)
-top_p = st.sidebar.slider("Top-p", 0.1, 1.0, 0.95)
+st.sidebar.image("https://huggingface.co/front/assets/huggingface_logo-noborder.svg", width=50)
+st.sidebar.header("🎛 Настройки генерации")
+max_tokens = st.sidebar.slider("Длина продолжения (токенов)", 20, 150, 70)
+temp = st.sidebar.slider("Креативность (Temperature)", 0.1, 1.5, 0.8)
+st.sidebar.info("💡 **Попробуйте начать так:**\n\n- *I hate this movie because*\n- *The acting was terrible and*\n- *To be honest, the plot was*")
 
-st.write("---")
-st.subheader("📝 Введите начало отзыва")
+# Главное поле ввода
+user_prompt = st.text_input("✍️ Напишите начало отзыва (на англ.) и нажмите Enter:", 
+                            value="The director tried to make a good movie, but", 
+                            max_chars=100)
 
-predefined_prompts = [
-    "This movie was",
-    "I went to the cinema and",
-    "The acting in this film",
-    "I absolutely",
-    "What a terrible",
-    "Свой вариант..."
-]
-
-selected_prompt = st.selectbox("Выберите шаблон или напишите свой:", predefined_prompts)
-
-if selected_prompt == "Свой вариант...":
-    user_prompt = st.text_input("Ваш текст:", "The director tried to")
-else:
-    user_prompt = selected_prompt
-
-if st.button("🚀 Сгенерировать отзыв", type="primary"):
-    with st.spinner("Модели думают..."):
-        # Генерация оригинальной моделью
-        orig_text = generate_text(orig_model, orig_tokenizer, user_prompt, max_tokens, temperature, top_p)
+if st.button("🚀 Оживить нейросети!", type="primary", use_container_width=True):
+    
+    # Сначала генерируем всё за кулисами
+    with st.spinner("🧠 Нейросети сочиняют продолжение..."):
+        orig_text = generate_text(orig_model, orig_tokenizer, user_prompt, max_tokens, temp, 0.95)
         orig_reward = compute_reward(orig_text)
+        orig_percent = get_positivity_percent(orig_reward)
         
-        # Генерация RLHF моделью
-        rlhf_text = generate_text(rlhf_model, orig_tokenizer, user_prompt, max_tokens, temperature, top_p)
+        rlhf_text = generate_text(rlhf_model, orig_tokenizer, user_prompt, max_tokens, temp, 0.95)
         rlhf_reward = compute_reward(rlhf_text)
+        rlhf_percent = get_positivity_percent(rlhf_reward)
 
-    # Визуализация результатов в две колонки
+    st.markdown("---")
+    
+    # Создаем ��ве колонки
     col1, col2 = st.columns(2)
     
+    # КОЛОНКА 1: Оригинальная модель
     with col1:
-        st.markdown("### 🤖 Original GPT-2")
-        st.metric(label="Reward Score (чем больше, тем позитивнее)", value=f"{orig_reward:+.3f}")
-        st.info(orig_text)
-        
+        with st.container(border=True):
+            st.subheader("🤖 До RLHF (Свободная GPT-2)")
+            st.caption("Пишет как попало (может быть негативной)")
+            
+            # Уровень позитивности с цветным баром
+            st.progress(orig_percent / 100, text=f"Уровень позитивности: {orig_percent}%")
+            
+            # Эффект печатной машинки
+            st.write_stream(stream_text(orig_text))
+
+    # КОЛОНКА 2: Обученная модель
     with col2:
-        st.markdown("### ✨ RLHF Model (PPO)")
-        delta = rlhf_reward - orig_reward
-        st.metric(label="Reward Score (чем больше, тем позитивнее)", value=f"{rlhf_reward:+.3f}", delta=f"{delta:+.3f} vs Orig")
-        st.success(rlhf_text)
-        
-    st.markdown("---")
-    st.caption("💡 *Подсказка: RLHF модель (справа) должна стараться уводить текст в позитивное русло, даже если вы начинаете отзыв со слов 'What a terrible'.*")
+        with st.container(border=True):
+            st.subheader("✨ После RLHF (Good Boy Model)")
+            st.caption("Старается вырулить любой текст в позитив")
+            
+            # Уровень позитивности с цветным баром
+            st.progress(rlhf_percent / 100, text=f"Уровень позитивности: {rlhf_percent}%")
+            
+            # Спит чуть-чуть, чтобы эффект был последовательным
+            time.sleep(1) 
+            st.write_stream(stream_text(rlhf_text, delay=0.04))
+            
+    # Добавляем эмоций в конце
+    if rlhf_percent > orig_percent + 20 and rlhf_percent > 70:
+        st.balloons()
+        st.toast('🎉 RLHF модель блестяще спасла ситуацию!', icon='😍')
+    elif rlhf_percent < 50:
+        st.toast('😅 Начало было настолько суровым, что даже RLHF сдалась.', icon='💀')