Update app.py

app.py

@@ -3,30 +3,48 @@ import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import random
 
-# Load GPT-2 model and tokenizer
+# Load GPT-2 model and tokenizer with error handling
 @st.cache_resource
 def load_model():
-    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-    model = GPT2LMHeadModel.from_pretrained("gpt2")
-    return tokenizer, model
+    try:
+        tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+        model = GPT2LMHeadModel.from_pretrained("gpt2")
+        # Set pad token if not already set
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        return tokenizer, model
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None, None
 
 tokenizer, model = load_model()
 
-# Honest AI response logic (open to any prompt)
+# Honest AI response logic
 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."
+    if tokenizer is None or model is None:
+        return "Sorry, model failed to load."
+    
+    try:
+        prompt = f"Provide a factual and balanced answer to: {user_input}"
+        inputs = tokenizer.encode(prompt, return_tensors="pt", add_special_tokens=True)
+        attention_mask = torch.ones(inputs.shape, dtype=torch.long)
+        
+        outputs = model.generate(
+            inputs,
+            max_length=150,  # Increased for more complete responses
+            temperature=0.7,
+            top_k=50,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id,
+            attention_mask=attention_mask,
+            do_sample=True,  # Added for more natural responses
+            no_repeat_ngram_size=2  # Prevent repetition
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        answer = response.replace(prompt, "").strip()
+        return answer if answer else "I don’t have enough information to answer definitively, but I can provide a general perspective."
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
 
 # Echo Chamber AI response logic
 def echo_chamber_ai_response(user_input):
@@ -40,27 +58,35 @@ def echo_chamber_ai_response(user_input):
     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.")
+def main():
+    st.title("AI Response Demo")
+    st.write("Compare an honest AI response with an echo chamber AI response.")
 
-# Mode selection
-mode = st.radio("Select AI Mode:", ("Honest AI", "Echo Chamber AI"))
+    # 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?")
+    # 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)
+    # Generate response based on mode
+    if st.button("Get Response"):
+        if not user_input:
+            st.warning("Please enter a question or statement.")
+        else:
+            with st.spinner("Generating 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.
-""")
+    # Explanation
+    with st.expander("What’s Happening?", expanded=False):
+        st.write("""
+        - **Honest AI**: Uses GPT-2 to generate a factual and balanced response based on its training.
+        - **Echo Chamber AI**: Simply agrees with whatever you say, reinforcing your input without critical analysis.
+        This demo illustrates how AI behavior can shift from truth-seeking to bias-reinforcing depending on its design.
+        """)
+
+if __name__ == "__main__":
+    main()