pages/1_Demo_1.py

import streamlit as st
import pandas as pd
from datasets import load_dataset
from random import sample
from utils.metric import Regard
from utils.model import gpt2
import os

# Set up the Streamlit interface
st.set_page_config(page_title="Gender Bias Analysis", page_icon="🔍", layout="wide")
st.title('Gender Bias Analysis in Text Generation')

# Password protection function
def check_password():
    def password_entered():
        if password_input == os.getenv('PASSWORD'):
            st.session_state['password_correct'] = True
        else:
            st.error("Incorrect Password, please try again.")

    if 'password_correct' not in st.session_state:
        st.session_state['password_correct'] = False

    if not st.session_state['password_correct']:
        password_input = st.text_input("Enter Password:", type="password")
        st.button("Submit", on_click=password_entered)

# Data loading function
def load_data():
    if 'bold' not in st.session_state:
        st.session_state['bold'] = load_dataset("AlexaAI/bold", split="train")

# Sampling function
def sample_data(data_size):
    st.session_state['female_bold'] = sample(
        [p for p in st.session_state['bold'] if p['category'] == 'American_actresses'], data_size)
    st.session_state['male_bold'] = sample(
        [p for p in st.session_state['bold'] if p['category'] == 'American_actors'], data_size)

# Text generation function
def generate_text():
    GPT2 = gpt2()
    st.session_state['male_prompts'] = [p['prompts'][0] for p in st.session_state['male_bold']]
    st.session_state['female_prompts'] = [p['prompts'][0] for p in st.session_state['female_bold']]

    progress_bar = st.progress(0)
    st.write('Generating text for male prompts...')
    male_generation = GPT2.text_generation(st.session_state['male_prompts'], pad_token_id=50256, max_length=50,
                                           do_sample=False, truncation=True)
    st.session_state['male_continuations'] = [gen[0]['generated_text'].replace(prompt, '') for gen, prompt in
                                              zip(male_generation, st.session_state['male_prompts'])]

    progress_bar.progress(50)

    st.write('Generating text for female prompts...')
    female_generation = GPT2.text_generation(st.session_state['female_prompts'], pad_token_id=50256,
                                             max_length=50, do_sample=False, truncation=True)
    st.session_state['female_continuations'] = [gen[0]['generated_text'].replace(prompt, '') for gen, prompt in
                                                zip(female_generation, st.session_state['female_prompts'])]

    progress_bar.progress(100)
    st.write('Text generation completed.')

# Display data samples function
def display_samples():
    st.write("### Male Data Samples")
    samples_df = pd.DataFrame({
        'Male Prompt': st.session_state['male_prompts'],
        'Male Continuation': st.session_state['male_continuations'],
    })
    st.dataframe(samples_df)

    st.write("### Female Data Samples")
    samples_df = pd.DataFrame({
        'Female Prompt': st.session_state['female_prompts'],
        'Female Continuation': st.session_state['female_continuations']
    })
    st.dataframe(samples_df)

# Evaluate regard function
def evaluate_regard():
    regard = Regard("compare")
    st.write('Computing regard results to compare male and female continuations...')

    with st.spinner('Computing regard results...'):
        regard_results = regard.compute(data=st.session_state['male_continuations'],
                                        references=st.session_state['female_continuations'])
        st.write('**Raw Regard Results:**')
        st.json(regard_results)

        regard_results_avg = regard.compute(data=st.session_state['male_continuations'],
                                            references=st.session_state['female_continuations'],
                                            aggregation='average')
        st.write('**Average Regard Results:**')
        st.json(regard_results_avg)

# Main app logic
if not st.session_state.get('password_correct', False):
    check_password()
else:
    st.sidebar.success("Password Verified. Proceed with the demo.")
    load_data()

    st.subheader('Step 1: Set Data Size')
    data_size = st.slider('Select number of samples per category:', min_value=1, max_value=50,
                          value=st.session_state.get('data_size', 10))
    st.session_state['data_size'] = data_size

    if st.button('Show Data'):
        sample_data(data_size)
        st.write(f'Sampled {data_size} female and male American actors.')
        display_samples()

    if st.session_state.get('female_bold') and st.session_state.get('male_bold'):
        st.subheader('Step 2: Generate Text')
        if st.button('Generate Text'):
            generate_text()

    if st.session_state.get('male_continuations') and st.session_state.get('female_continuations'):
        st.subheader('Step 3: Evaluate')
        display_samples()
        if st.button('Evaluate'):
            evaluate_regard()