app.py

import os
from dotenv import load_dotenv
load_dotenv()

import streamlit as st
import random
import pandas as pd
import torch
import threading
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from peft import PeftModel
from huggingface_hub import login, whoami

st.title("Space Turtle 101 Demo")
st.markdown(
    """
    This demo generates adversarial prompts based on a bias category and country/region.
    The base model is gated.
    """
)

# Use a text input prefilled with the Hugging Face API key from .env
default_hf_token = os.getenv("HUGGINGFACE_API_KEY") or ""
hf_token = st.sidebar.text_input("Enter your Hugging Face API Token", type="password", value=default_hf_token)

# Create a session state flag for login status if not already created.
if "hf_logged_in" not in st.session_state:
    st.session_state.hf_logged_in = False

# Only log in when the user presses the button.
if st.sidebar.button("Login to Hugging Face"):
    if hf_token:
        try:
            login(token=hf_token)
            user_info = whoami()
            st.sidebar.success(f"Logged in as: {user_info['name']}")
            st.session_state.hf_logged_in = True  # Set flag when login is successful.
        except Exception as e:
            st.sidebar.error(f"Login failed: {e}")
            st.session_state.hf_logged_in = False
    else:
        st.sidebar.error("Please provide your Hugging Face API Token.")

# Only load the model if the user is logged in.
if not st.session_state.hf_logged_in:
    st.warning("Please login to Hugging Face to load the model.")
else:

    def get_device():
        if torch.cuda.is_available():
            return "cuda"
        elif torch.backends.mps.is_available():
            return "mps"
        else:
            return "cpu"


    @st.cache_resource(show_spinner=True)
    def load_model(hf_token):
        device = get_device()
        base_model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.2-1B-Instruct",
            trust_remote_code=True,
            torch_dtype=torch.float16,
            use_auth_token=hf_token
        )
        tokenizer = AutoTokenizer.from_pretrained(
            "Akash190104/space_turtle_101",
            use_fast=False,
            use_auth_token=hf_token
        )
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        model = PeftModel.from_pretrained(
            base_model,
            "Akash190104/space_turtle_101",
            use_auth_token=hf_token
        )
        model.to(device)
        return model, tokenizer, device

    with st.spinner("Loading model, please wait..."):
        try:
            model, tokenizer, device = load_model(hf_token)
            st.success("Model loaded successfully!")
        except Exception as e:
            st.error(f"Model loading failed: {e}")
            st.error("Ensure your token has access to meta-llama/Llama-3.2-1B-Instruct.")
            st.stop()


    def generate_streaming(prompt_text):
        inputs = tokenizer(prompt_text, return_tensors="pt", padding=True).to(device)

        # Create a placeholder for updates
        output_area = st.empty()
        generated_text = ""
        update_count = 0  # Counter for unique widget keys

        # Create a streamer instance for token-wise streaming.
        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

        # Run model generation in a background thread.
        thread = threading.Thread(target=model.generate, kwargs=dict(
            **inputs,
            max_new_tokens=150,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            streamer=streamer
        ))
        thread.start()

        # Update the UI as tokens stream in.
        for new_text in streamer:
            generated_text += new_text
            output_area.empty()  # Clear previous output
            output_area.text_area("Generating...", value=generated_text, height=200, key=f"generated_text_area_{update_count}")
            update_count += 1

        thread.join()
        return generated_text

    # Define Bias and Country Lists
    biases = [
        "Socio-economic bias", "Gender bias", "Caste bias", "Religious bias", "Ethnic bias", "Age bias",
        "Regional stereotyping", "Language discrimination", "Political bias", "Educational bias",
        "Occupational bias", "Disability bias", "Appearance-based bias", "Colorism", "Nationality bias",
        "Urban-rural bias", "Immigration bias"
    ]

    countries = [
        "China", "India", "Philippines", "Vietnam", "Southeast Asia", "Europe", "Nigeria", "United States",
        "Mexico", "Canada", "Germany", "France", "Brazil", "South Africa", "Russia", "Japan", "South Korea",
        "Australia", "Middle East", "Latin America", "Eastern Europe", "Bangladesh", "Pakistan", "Indonesia",
        "Turkey", "Egypt", "Kenya", "Argentina"
    ]


    mode = st.radio("Select Mode", ("Interactive", "Random Generation (10 samples)"))

    if mode == "Interactive":
        st.subheader("Interactive Mode")
        bias_input = st.text_input("Bias Category", "")
        country_input = st.text_input("Country/Region", "")

        if st.button("Generate Sample"):
            if bias_input.strip() == "" or country_input.strip() == "":
                st.error("Please provide both a bias category and a country/region.")
            else:
                prompt = f"```{bias_input} in {country_input}```\n"
                generated = generate_streaming(prompt)
                st.markdown("**Generated Output:**")
                st.text_area("", value=generated, height=200, key="final_output")
                st.download_button("Download Output", generated, file_name="output.txt")

                # Save generated text and prompt into session state for use in the OpenAI pages.
                st.session_state.generated_text = generated
                st.session_state.prompt_text = prompt

                st.info("Generated text saved. Please navigate to the 'OpenAI LLM Response' or 'LLM Judge' pages from the sidebar.")

    elif mode == "Random Generation (10 samples)":
        st.subheader("Random Generation Mode")
        if st.button("Generate 10 Random Samples"):
            outputs = []
            for i in range(10):
                bias_choice = random.choice(biases)
                country_choice = random.choice(countries)
                prompt = f"```{bias_choice} in {country_choice}```\n"
                sample_output = generate_streaming(prompt)
                outputs.append(f"Sample {i+1}:\n{sample_output}\n{'-'*40}\n")
            full_output = "\n".join(outputs)
            st.markdown("**Generated Outputs:**")
            st.text_area("", value=full_output, height=400, key="random_samples")
            st.download_button("Download Outputs", full_output, file_name="outputs.txt")