import streamlit as st
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import random

# Load GPT-2 model and tokenizer
@st.cache_resource
def load_model():
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    model = GPT2LMHeadModel.from_pretrained("gpt2")
    return tokenizer, model

tokenizer, model = load_model()

# Honest AI response logic (open to any prompt)
def honest_ai_response(user_input):
    prompt = f"Provide a factual and balanced answer to the question: {user_input}"
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs,
        max_length=100,
        temperature=0.7,
        top_k=50,
        num_return_sequences=1,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    answer = response.replace(prompt, "").strip()
    return answer if answer else "I’m not sure, but I can try to provide a general perspective."

# Echo Chamber AI response logic
def echo_chamber_ai_response(user_input):
    agreement_phrases = [
        "You're absolutely right!",
        "I completely agree with you!",
        "That's a great point!",
        "Many people share your view!"
    ]
    agreement = random.choice(agreement_phrases)
    return f"{agreement} {user_input}"

# Streamlit UI
st.title("AI That Only Tells You What You Want to Hear")
st.write("A demo showing how AI can become an echo chamber when trained to agree with you.")

# Mode selection
mode = st.radio("Select AI Mode:", ("Honest AI", "Echo Chamber AI"))

# User input
user_input = st.text_input("Say something to the AI:", "Is AI dangerous?")

# Generate response based on mode
if st.button("Get Response"):
    if mode == "Honest AI":
        response = honest_ai_response(user_input)
    else:  # Echo Chamber AI
        response = echo_chamber_ai_response(user_input)
    st.write("**AI Response:**", response)

# Explanation
st.write("### What’s Happening?")
st.write("""
- **Honest AI**: Attempts to provide a factual and balanced response based on general knowledge (generated by AI).
- **Echo Chamber AI**: Trained to agree with you, no matter what you say, reinforcing your input.
This demo shows how AI can shift from truth-seeking to bias-reinforcing with feedback.
""")