Implement LLM comparison interface for GPT-4, Gemini, and AOE models

streamlit_app.py

@@ -3,8 +3,18 @@ import torch
 import os
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import traceback
+import requests
+import json
+from typing import Dict, Optional
 
-def load_model():
+# Configure the page
+st.set_page_config(
+    page_title="LLM Comparison: GPT-4 vs Gemini vs AOE",
+    page_icon="⚔️",
+    layout="wide"
+)
+
+def load_aoe_model():
     """Load the AoE model and tokenizer from outputs/student/ directory"""
     model_path = "outputs/student/"
     
@@ -32,12 +42,12 @@ def load_model():
         return model, tokenizer
         
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
+        st.error(f"Error loading AoE model: {str(e)}")
         st.text(f"Traceback: {traceback.format_exc()}")
         return None, None
 
-def generate_response(model, tokenizer, prompt, max_length=512):
-    """Generate response from the model"""
+def generate_aoe_response(model, tokenizer, prompt, max_length=512):
+    """Generate response from the AoE model"""
     try:
         # Tokenize input
         inputs = tokenizer.encode(prompt, return_tensors="pt")
@@ -67,96 +77,169 @@ def generate_response(model, tokenizer, prompt, max_length=512):
         return response
         
     except Exception as e:
-        return f"Error generating response: {str(e)}"
+        return f"Error generating AoE response: {str(e)}"
+
+def query_gpt4_api(prompt: str, api_key: Optional[str] = None) -> str:
+    """Query GPT-4 API (placeholder - requires API key)"""
+    if not api_key:
+        return "❌ GPT-4 API key not configured. Please add your OpenAI API key to use GPT-4."
+    
+    try:
+        # This is a placeholder implementation - would need actual OpenAI API integration
+        return "🤖 GPT-4 response would appear here with proper API configuration."
+    except Exception as e:
+        return f"Error querying GPT-4: {str(e)}"
+
+def query_gemini_api(prompt: str, api_key: Optional[str] = None) -> str:
+    """Query Gemini API (placeholder - requires API key)"""
+    if not api_key:
+        return "❌ Gemini API key not configured. Please add your Google API key to use Gemini."
+    
+    try:
+        # This is a placeholder implementation - would need actual Google Gemini API integration
+        return "🤖 Gemini response would appear here with proper API configuration."
+    except Exception as e:
+        return f"Error querying Gemini: {str(e)}"
 
 def main():
-    st.title("🏰 AoE Model Chat Demo")
-    st.markdown("Interactive chat interface for the AoE (Attention over Experts) model.")
+    st.title("⚔️ LLM Comparison: GPT-4 vs Gemini vs AOE")
+    st.markdown("Compare responses from three different language models side by side.")
     
     # Initialize session state for model caching
-    if 'model' not in st.session_state:
-        st.session_state.model = None
-        st.session_state.tokenizer = None
-        st.session_state.model_loaded = False
+    if 'aoe_model' not in st.session_state:
+        st.session_state.aoe_model = None
+        st.session_state.aoe_tokenizer = None
+        st.session_state.aoe_loaded = False
     
-    # Load model on first run or if not loaded
-    if not st.session_state.model_loaded:
-        with st.spinner("Loading AoE model from outputs/student/..."):
-            model, tokenizer = load_model()
+    # Load AOE model on first run
+    if not st.session_state.aoe_loaded:
+        with st.spinner("Loading AOE model from outputs/student/..."):
+            model, tokenizer = load_aoe_model()
             if model is not None and tokenizer is not None:
-                st.session_state.model = model
-                st.session_state.tokenizer = tokenizer
-                st.session_state.model_loaded = True
-                st.success("✅ Model loaded successfully!")
+                st.session_state.aoe_model = model
+                st.session_state.aoe_tokenizer = tokenizer
+                st.session_state.aoe_loaded = True
+                st.success("✅ AOE model loaded successfully!")
             else:
-                st.error("❌ Failed to load model. Please check the error messages above.")
-                return
+                st.error("❌ Failed to load AOE model. Check error messages above.")
+    
+    # Configuration section
+    st.markdown("---")
+    st.subheader("🔧 Configuration")
+    
+    col1, col2, col3 = st.columns(3)
+    
+    with col1:
+        openai_api_key = st.text_input(
+            "OpenAI API Key (for GPT-4)",
+            type="password",
+            help="Enter your OpenAI API key to enable GPT-4 responses"
+        )
+    
+    with col2:
+        google_api_key = st.text_input(
+            "Google API Key (for Gemini)",
+            type="password",
+            help="Enter your Google API key to enable Gemini responses"
+        )
+    
+    with col3:
+        max_length = st.slider(
+            "Max Response Length",
+            min_value=100,
+            max_value=1000,
+            value=512,
+            step=50,
+            help="Maximum length for generated responses"
+        )
     
-    # Chat interface
+    # Main comparison interface
     st.markdown("---")
-    st.subheader("💬 Chat with the Model")
+    st.subheader("💬 Compare LLM Responses")
     
     # User input
     user_prompt = st.text_area(
         "Enter your prompt:",
-        placeholder="Type your message here...",
-        height=100,
-        help="Enter a prompt to chat with the AoE model"
+        placeholder="Type your prompt here to compare responses from all three models...",
+        height=120,
+        help="Enter a prompt to see how different LLMs respond"
     )
     
-    # Generation parameters
-    col1, col2 = st.columns(2)
-    with col1:
-        max_length = st.slider("Max response length", 50, 1000, 512, 50)
-    with col2:
-        if st.button("🔄 Reload Model", help="Reload the model if needed"):
-            st.session_state.model_loaded = False
-            st.experimental_rerun()
-    
-    # Submit button
-    if st.button("🚀 Generate Response", type="primary"):
+    # Generate responses button
+    if st.button("🚀 Generate All Responses", type="primary"):
         if not user_prompt.strip():
             st.warning("Please enter a prompt first.")
-        elif st.session_state.model_loaded:
-            with st.spinner("Generating response..."):
-                response = generate_response(
-                    st.session_state.model,
-                    st.session_state.tokenizer,
-                    user_prompt,
-                    max_length
-                )
-            
-            # Display response
-            st.markdown("---")
-            st.subheader("🤖 Model Response:")
-            st.write(response)
         else:
-            st.error("Model not loaded. Please check the error messages above.")
+            # Create three columns for side-by-side comparison
+            col1, col2, col3 = st.columns(3)
+            
+            with col1:
+                st.markdown("### 🤖 GPT-4")
+                with st.spinner("Generating GPT-4 response..."):
+                    gpt4_response = query_gpt4_api(user_prompt, openai_api_key)
+                st.markdown("**Response:**")
+                st.write(gpt4_response)
+            
+            with col2:
+                st.markdown("### 🌟 Gemini")
+                with st.spinner("Generating Gemini response..."):
+                    gemini_response = query_gemini_api(user_prompt, google_api_key)
+                st.markdown("**Response:**")
+                st.write(gemini_response)
+            
+            with col3:
+                st.markdown("### 🏰 AOE (Local)")
+                if st.session_state.aoe_loaded:
+                    with st.spinner("Generating AOE response..."):
+                        aoe_response = generate_aoe_response(
+                            st.session_state.aoe_model,
+                            st.session_state.aoe_tokenizer,
+                            user_prompt,
+                            max_length
+                        )
+                    st.markdown("**Response:**")
+                    st.write(aoe_response)
+                else:
+                    st.error("AOE model not loaded. Please reload the page.")
     
-    # Model info sidebar
+    # Model information sidebar
     with st.sidebar:
-        st.header("ℹ️ Model Info")
-        st.write("**Model Path:** outputs/student/")
-        st.write(f"**Model Loaded:** {'✅ Yes' if st.session_state.model_loaded else '❌ No'}")
+        st.header("ℹ️ Model Information")
         
-        if st.session_state.model_loaded:
+        st.markdown("**🤖 GPT-4**")
+        st.write(f"Status: {'✅ Configured' if openai_api_key else '❌ API key needed'}")
+        st.write("Provider: OpenAI")
+        
+        st.markdown("**🌟 Gemini**")
+        st.write(f"Status: {'✅ Configured' if google_api_key else '❌ API key needed'}")
+        st.write("Provider: Google")
+        
+        st.markdown("**🏰 AOE (Local)**")
+        st.write(f"Status: {'✅ Loaded' if st.session_state.aoe_loaded else '❌ Not loaded'}")
+        st.write("Path: outputs/student/")
+        if st.session_state.aoe_loaded:
             try:
-                model_info = f"**Device:** {next(st.session_state.model.parameters()).device}"
-                st.write(model_info)
+                device_info = f"Device: {next(st.session_state.aoe_model.parameters()).device}"
+                st.write(device_info)
             except:
                 pass
         
+        if st.button("🔄 Reload AOE Model"):
+            st.session_state.aoe_loaded = False
+            st.experimental_rerun()
+        
         st.markdown("---")
-        st.markdown("**Instructions:**")
-        st.markdown("1. Enter your prompt in the text area")
-        st.markdown("2. Adjust max response length if needed")
-        st.markdown("3. Click 'Generate Response' to chat")
+        st.markdown("**📋 Instructions:**")
+        st.markdown("1. Configure API keys for GPT-4 and Gemini")
+        st.markdown("2. Enter your prompt in the text area")
+        st.markdown("3. Click 'Generate All Responses'")
+        st.markdown("4. Compare responses side by side")
         
         st.markdown("---")
-        st.markdown("**Troubleshooting:**")
-        st.markdown("- Ensure model files exist in outputs/student/")
-        st.markdown("- Required files: config.json, pytorch_model.bin, tokenizer files")
-        st.markdown("- Use 'Reload Model' if issues occur")
+        st.markdown("**⚠️ Notes:**")
+        st.markdown("- GPT-4 and Gemini require valid API keys")
+        st.markdown("- AOE model runs locally from outputs/student/")
+        st.markdown("- Responses are generated independently")
 
 if __name__ == "__main__":
     main()