app.py

import streamlit as st
import random
from config import initialize
from firebase_utils import load_ratings, save_vote
from openai_utils import generate_rewrite
from data_utils import load_data, get_random_review, read_prompt

models = ["gpt-3.5-turbo", "gpt-4", "gpt-4o", "gpt-4o-mini"]

def local_css(file_name):
    with open(file_name) as f:
        css = f.read()
        st.markdown(f"<style>{css}</style>", unsafe_allow_html=True)

def main():
    st.set_page_config(page_title="AI Арена", layout="wide")
    local_css("styles.css")  
    st.markdown("<h1 style='text-align: center;'>⚔️ AI Арена - Битва Моделей</h1>", unsafe_allow_html=True)

    st.markdown("<h3>🎲 Нажми кнопку ниже, чтобы загрузить оригинальный отзыв и начать битву!</h3>", unsafe_allow_html=True)

    db = initialize()

    data = load_data('final_data_wb.xlsx')

    if 'original_review' not in st.session_state:
        st.session_state.original_review = ''
    if 'model_a' not in st.session_state:
        st.session_state.model_a = ''
    if 'model_b' not in st.session_state:
        st.session_state.model_b = ''
    if 'rewrite_a' not in st.session_state:
        st.session_state.rewrite_a = ''
    if 'rewrite_b' not in st.session_state:
        st.session_state.rewrite_b = ''
    if 'vote_submitted' not in st.session_state:
        st.session_state.vote_submitted = False

    instruction = read_prompt('prompt.txt')

    if st.button("🚀 Начать битву"):
        st.session_state.original_review = get_random_review(data)

        selected_models = random.sample(models, 2)
        st.session_state.model_a = selected_models[0]
        st.session_state.model_b = selected_models[1]
        st.session_state.rewrite_a = ''
        st.session_state.rewrite_b = ''
        st.session_state.vote_submitted = False

    if st.session_state.original_review:
        st.markdown("<h2>📝 Оригинальный отзыв:</h2>", unsafe_allow_html=True)
        st.markdown(
            f"<div class='original-review'>{st.session_state.original_review}</div>",
            unsafe_allow_html=True
        )

        if st.button("⚔️ Сгенерировать рерайты"):
            with st.spinner("Генерация рерайтов..."):
                st.session_state.rewrite_a = generate_rewrite(
                    st.session_state.model_a, instruction, st.session_state.original_review)
                st.session_state.rewrite_b = generate_rewrite(
                    st.session_state.model_b, instruction, st.session_state.original_review)

    if st.session_state.rewrite_a and st.session_state.rewrite_b:
        st.markdown("""
            <h2>🗳️ Проголосуй за лучший рерайт!</h2>
            <p>Выбери рерайт, который соответствует данным критериям:</p>
            <ol>
                <li>Сохранен смысл оригинального отзыва, но написан другими словами.</li>
                <li>Написан повседневным языком и не похож на сгенерированный.</li>
                <li>Имеет структуру <strong>Отзыв</strong>, <strong>Преимущества</strong>, <strong>Недостатки</strong>.</li>
            </ol>
        """, unsafe_allow_html=True)
        cols = st.columns(2)
        with cols[0]:
            st.markdown("<h3>Model A</h3>", unsafe_allow_html=True)
            st.markdown(
            f"<div class='original-review'>{st.session_state.rewrite_a}</div>",
            unsafe_allow_html=True
            )
            st.markdown("</div>", unsafe_allow_html=True)
            if not st.session_state.vote_submitted:
                if st.button("👍 Выбрать Model A"):
                    st.session_state.vote_submitted = True
                    st.success("Вы выбрали Model A")
                    save_vote(db, 'left', models)
        with cols[1]:
            st.markdown("<h3>Model B</h3>", unsafe_allow_html=True)
            st.markdown(
            f"<div class='original-review'>{st.session_state.rewrite_b}</div>",
            unsafe_allow_html=True
            )
            st.markdown("</div>", unsafe_allow_html=True)
            if not st.session_state.vote_submitted:
                if st.button("👍 Выбрать Model B"):
                    st.session_state.vote_submitted = True
                    st.success("Вы выбрали Model B")
                    save_vote(db, 'right', models)

if __name__ == '__main__':
    main()