Update app.py

app.py

@@ -5,411 +5,446 @@ import threading
 import torch
 import os
 import time
-from typing import List, Dict, Generator, Tuple, Optional
+from typing import List, Dict, Generator, Tuple, Optional, Union
 import logging
+import warnings
+from dataclasses import dataclass
+import gc
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# --- Best Free Models for Council ---
+warnings.filterwarnings("ignore", message="torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly")
+
+@dataclass
+class ModelInfo:
+    id: str
+    name: str
+    required_memory: str  # Estimated VRAM requirement
+
+@dataclass
+class Persona:
+    name: str
+    description: str
+    traits: str
+    style: str
+    emoji: str
+
 MODELS = [
-    ("meta-llama/Meta-Llama-3-8B-Instruct", "Llama 3 8B Instruct"),
-    ("Qwen/Qwen1.5-7B-Chat", "Qwen1.5 7B Chat"),
-    ("HuggingFaceH4/zephyr-7b-beta", "Zephyr 7B Beta"),
-    ("mistralai/Mistral-7B-Instruct-v0.2", "Mistral 7B Instruct"),
+    ModelInfo("meta-llama/Meta-Llama-3-8B-Instruct", "Llama 3 8B Instruct", "16GB"),
+    ModelInfo("Qwen/Qwen1.5-7B-Chat", "Qwen1.5 7B Chat", "14GB"),
+    ModelInfo("HuggingFaceH4/zephyr-7b-beta", "Zephyr 7B Beta", "14GB"),
+    ModelInfo("mistralai/Mistral-7B-Instruct-v0.2", "Mistral 7B Instruct", "14GB"),
 ]
 
-# Define council member personas with enhanced characteristics
 PERSONAS = [
-    {
-        "name": "Dr. Ana Rodriguez",
-        "description": "An analytical scientist who values empirical evidence and logical reasoning. Often plays devil's advocate and questions assumptions.",
-        "traits": "analytical, skeptical, evidence-focused",
-        "style": "formal, precise, methodical",
-        "emoji": "🔬"
-    },
-    {
-        "name": "Professor Marcus Chen",
-        "description": "A creative philosopher with an interest in ethics and societal implications. Considers the bigger picture and long-term consequences.",
-        "traits": "philosophical, visionary, empathetic",
-        "style": "eloquent, metaphorical, conceptual",
-        "emoji": "🧠"
-    },
-    {
-        "name": "Sarah Johnson",
-        "description": "A pragmatic problem-solver with real-world experience. Focuses on practicality and implementation details.",
-        "traits": "practical, solution-oriented, experienced",
-        "style": "direct, concise, example-driven",
-        "emoji": "🛠️"
-    },
-    {
-        "name": "Dr. Emeka Okafor",
-        "description": "A social scientist specializing in cultural perspectives and community impacts. Brings diverse viewpoints and contextual understanding.",
-        "traits": "culturally aware, nuanced, community-focused",
-        "style": "inclusive, storytelling, perspective-oriented",
-        "emoji": "🌍"
-    }
+    Persona(
+        name="Dr. Ana Rodriguez",
+        description="An analytical scientist who values empirical evidence and logical reasoning.",
+        traits="analytical, skeptical, evidence-focused",
+        style="formal, precise, methodical",
+        emoji="🔬"
+    ),
+    Persona(
+        name="Professor Marcus Chen",
+        description="A creative philosopher with an interest in ethics and societal implications.",
+        traits="philosophical, visionary, empathetic",
+        style="eloquent, metaphorical, conceptual",
+        emoji="🧠"
+    ),
+    Persona(
+        name="Sarah Johnson",
+        description="A pragmatic problem-solver with real-world experience.",
+        traits="practical, solution-oriented, experienced",
+        style="direct, concise, example-driven",
+        emoji="🛠️"
+    ),
+    Persona(
+        name="Dr. Emeka Okafor",
+        description="A social scientist specializing in cultural perspectives.",
+        traits="culturally aware, nuanced, community-focused",
+        style="inclusive, storytelling, perspective-oriented",
+        emoji="🌍"
+    )
 ]
 
-# Cache for models to avoid reloading
 model_cache = {}
+current_device = None
+
+def get_device() -> str:
+    global current_device
+    if current_device:
+        return current_device
+    if torch.cuda.is_available():
+        gpu_mem = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+        current_device = "cuda"
+        logger.info(f"GPU available with {gpu_mem:.1f}GB memory")
+        min_required = min(float(model.required_memory.replace("GB", "")) for model in MODELS)
+        if gpu_mem < min_required:
+            logger.warning(f"GPU memory may be insufficient for some models (has {gpu_mem:.1f}GB, needs {min_required}GB)")
+    else:
+        current_device = "cpu"
+        logger.info("Using CPU")
+    return current_device
+
+def clear_model_cache():
+    global model_cache
+    for model_id in list(model_cache.keys()):
+        del model_cache[model_id]
+    gc.collect()
+    torch.cuda.empty_cache()
+    model_cache = {}
+    logger.info("Model cache cleared")
 
-def load_model(model_id: str) -> Tuple[pipeline, AutoTokenizer]:
-    """Load model and tokenizer with caching to improve performance"""
+def load_model(model_info: ModelInfo) -> Tuple[pipeline, AutoTokenizer]:
     global model_cache
-    
-    if model_id in model_cache:
-        logger.info(f"Using cached model: {model_id}")
-        return model_cache[model_id]
-    
-    logger.info(f"Loading model: {model_id}")
+    if model_info.id in model_cache:
+        logger.info(f"Using cached model: {model_info.name}")
+        return model_cache[model_info.id]
+    device = get_device()
+    if device == "cuda":
+        gpu_mem = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+        required_mem = float(model_info.required_memory.replace("GB", ""))
+        if gpu_mem < required_mem:
+            logger.warning(f"Insufficient GPU memory for {model_info.name} (needs {required_mem}GB, has {gpu_mem:.1f}GB)")
+    logger.info(f"Loading {model_info.name} on {device}")
     try:
-        # Set environmental variables for optimizations
-        os.environ["TOKENIZERS_PARALLELISM"] = "true"
-        
-        # Load tokenizer and model
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        
-        # Determine if CUDA is available and set appropriate device
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        
-        # Configure model loading for memory efficiency
+        start_time = time.time()
+        tokenizer = AutoTokenizer.from_pretrained(model_info.id)
         model_kwargs = {
             "trust_remote_code": True,
-            "device_map": "auto",
+            "device_map": "auto" if device == "cuda" else None,
             "torch_dtype": torch.float16 if device == "cuda" else torch.float32
         }
-        
-        model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
-        
-        # Create pipeline with appropriate settings
-        pipe = pipeline("text-generation", 
-                       model=model, 
-                       tokenizer=tokenizer, 
-                       max_new_tokens=512, 
-                       device=model.device)
-        
-        # Cache the model and tokenizer
-        model_cache[model_id] = (pipe, tokenizer)
-        logger.info(f"Model loaded successfully: {model_id} on {device}")
+        with gr.Progress() as progress:
+            progress(0, desc=f"Loading {model_info.name}")
+            model = AutoModelForCausalLM.from_pretrained(model_info.id, **model_kwargs)
+            if device == "cuda":
+                model = model.to(device)
+            pipe = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                device=model.device
+            )
+        model_cache[model_info.id] = (pipe, tokenizer)
+        load_time = time.time() - start_time
+        logger.info(f"Loaded {model_info.name} in {load_time:.1f}s")
         return pipe, tokenizer
-    
     except Exception as e:
-        logger.error(f"Failed to load model {model_id}: {str(e)}")
+        logger.error(f"Failed to load {model_info.name}: {str(e)}")
         raise
 
-def create_debate_prompt(user_prompt: str, 
-                        persona: Dict, 
-                        debate_style: str = "Balanced", 
-                        previous_responses: Optional[List[str]] = None) -> str:
-    """Create a prompt that encourages a natural debate-like response with adjustable style"""
-    persona_desc = f"You are {persona['name']}, {persona['description']} Your communication style is {persona['style']}."
-    
-    # Adjust prompt based on debate style
-    style_guidance = ""
-    if debate_style == "Collaborative":
-        style_guidance = "Focus on building upon and synthesizing the ideas of others. Look for common ground and areas of agreement."
-    elif debate_style == "Adversarial":
-        style_guidance = "Challenge assumptions and present contrasting viewpoints. Don't be afraid to disagree strongly with others."
-    else:  # Balanced
-        style_guidance = "Present your authentic perspective while being respectful of other viewpoints. Balance critique with constructive ideas."
-    
-    if not previous_responses:
-        prompt = f"""{persona_desc}
+def create_debate_prompt(
+    user_prompt: str,
+    persona: Persona,
+    debate_style: str = "Balanced",
+    previous_responses: Optional[List[str]] = None
+) -> str:
+    style_guidance = {
+        "Collaborative": "Focus on building upon ideas and finding common ground.",
+        "Adversarial": "Challenge assumptions and present strong contrasting views.",
+        "Balanced": "Present your perspective while respecting others."
+    }.get(debate_style, "Present your authentic perspective.")
+    base_prompt = f"""You are {persona.name}, {persona.description}
+Your communication style: {persona.style}
+Traits: {persona.traits}
 
-You are part of a council debating the following topic:
-"{user_prompt}"
+You're in a council debating: "{user_prompt}"
 
 {style_guidance}
-
-Give your authentic perspective on this topic based on your persona. Be natural and conversational.
-Directly address the topic without hedging or being overly formal. Make specific points that others can respond to.
-Keep your response to 3-4 paragraphs maximum.
-
-{persona['name']}:"""
-    else:
+Respond naturally in 3-4 paragraphs."""
+    if previous_responses:
         debate_history = "\n\n".join(previous_responses)
-        prompt = f"""{persona_desc}
-
-You are part of a council debating the following topic:
-"{user_prompt}"
+        return f"""{base_prompt}
 
-{style_guidance}
-
-The debate so far:
+Current discussion:
 {debate_history}
 
-Now it's your turn to speak. Based on your persona and the previous speakers:
-- You may agree or disagree with previous points
-- Add new perspectives they missed
-- Point out flaws in reasoning or suggest compromises
-- Address someone directly if appropriate
-- Be authentic to your character - don't just summarize
+Now respond thoughtfully to the ongoing debate:
+{persona.name}:"""
+    return f"""{base_prompt}
 
-Give your natural, conversational response as if in a real discussion.
-Keep your response to 3-4 paragraphs maximum.
-
-{persona['name']}:"""
-    
-    return prompt
+Begin your response:
+{persona.name}:"""
 
 def create_synthesis_prompt(user_prompt: str, all_responses: List[str]) -> str:
-    """Create a prompt for the facilitator to synthesize the debate"""
     debate_history = "\n\n".join(all_responses)
-    prompt = f"""You are the Facilitator, responsible for synthesizing the council's discussion on:
-"{user_prompt}"
+    return f"""As the Facilitator, synthesize this discussion:
+Topic: "{user_prompt}"
 
-The full debate:
+Debate:
 {debate_history}
 
-Provide a thoughtful synthesis that:
-1. Identifies the key points of agreement and disagreement
-2. Highlights the most compelling insights from each perspective
-3. Draws a balanced conclusion that respects the nuance of the discussion
-4. Offers a path forward or recommendation when appropriate
-
-Be concise but comprehensive. Focus on substance over style.
-Keep your synthesis to 3-5 paragraphs maximum.
+Provide:
+1. Key agreements/disagreements
+2. Important insights
+3. Balanced conclusion
+4. Recommended next steps
 
 Facilitator:"""
-    return prompt
 
-def stream_model_response(pipe: pipeline, 
-                         tokenizer: AutoTokenizer,
-                         prompt: str, 
-                         speaker_name: str,
-                         temperature: float = 0.7) -> Generator[str, None, None]:
-    """Stream model responses with better error handling"""
+def stream_response(
+    pipe: pipeline,
+    tokenizer: AutoTokenizer,
+    prompt: str,
+    speaker_name: Optional[str] = None,
+    temperature: float = 0.7,
+    max_tokens: int = 512
+) -> Generator[str, None, None]:
     try:
-        # Set up the streamer
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
         input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(pipe.model.device)
-        
-        # Run model generation in a separate thread
         generation_kwargs = dict(
             input_ids=input_ids,
             streamer=streamer,
-            max_new_tokens=512,
+            max_new_tokens=max_tokens,
             do_sample=True,
             temperature=temperature,
             top_p=0.95,
             repetition_penalty=1.1,
             eos_token_id=tokenizer.eos_token_id,
         )
-        
-        thread = threading.Thread(
-            target=pipe.model.generate,
-            kwargs=generation_kwargs
-        )
+        thread = threading.Thread(target=pipe.model.generate, kwargs=generation_kwargs)
         thread.start()
-        
-        # Stream the response as it's generated
-        response = ""
+        buffer = ""
         for new_text in streamer:
-            response += new_text
-            # Add the emoji to the speaker name
-            yield f"**{speaker_name}:** {response.strip()}"
-        
+            buffer += new_text
+            if new_text and new_text[-1] in " .,;!?\n":
+                if speaker_name:
+                    yield f"**{speaker_name}:** {buffer.strip()}"
+                else:
+                    yield buffer.strip()
         thread.join()
-        return response.strip()
-    
+        if buffer.strip():
+            if speaker_name:
+                yield f"**{speaker_name}:** {buffer.strip()}"
+            else:
+                yield buffer.strip()
     except Exception as e:
-        logger.error(f"Error streaming response: {str(e)}")
-        yield f"**{speaker_name}:** [Error generating response. Please try again.]"
+        logger.error(f"Streaming error: {str(e)}")
+        error_msg = f"[Error: {str(e)}]"
+        yield f"**{speaker_name}:** {error_msg}" if speaker_name else error_msg
 
-def council_chat_stream(user_prompt: str, 
-                       num_members: int = 3, 
-                       debate_style: str = "Balanced",
-                       temperature: float = 0.7) -> Generator[str, None, None]:
-    """Generate a council debate with configurable number of members and style"""
-    # Validate inputs
+def council_chat_stream(
+    user_prompt: str,
+    num_members: int = 3,
+    debate_style: str = "Balanced",
+    temperature: float = 0.7,
+    selected_models: Optional[List[str]] = None,
+    continue_debate: bool = False,
+    history: Optional[List[str]] = None
+) -> Generator[str, None, None]:
     if not user_prompt.strip():
-        yield "Please enter a topic for the council to debate."
+        yield "Please enter a topic for debate."
         return
-    
+    num_members = max(2, min(num_members, len(PERSONAS)))
+    temperature = max(0.1, min(temperature, 1.0))
     start_time = time.time()
-    
-    # Determine which personas and models to use
-    selected_personas = random.sample(PERSONAS, min(num_members, len(PERSONAS)))
-    selected_models = random.sample(MODELS, min(num_members, len(MODELS)))
-    
-    # Load models
+    selected_personas = random.sample(PERSONAS, num_members)
+    model_pool = selected_models if selected_models else [model.id for model in MODELS]
+    selected_model_infos = random.sample([m for m in MODELS if m.id in model_pool], num_members)
     loaded_models = []
-    for model_id, _ in selected_models:
+    for model_info in selected_model_infos:
         try:
-            pipe, tokenizer = load_model(model_id)
-            loaded_models.append((pipe, tokenizer))
+            with gr.Progress() as progress:
+                progress(0, desc=f"Loading {model_info.name}")
+                pipe, tokenizer = load_model(model_info)
+                loaded_models.append((pipe, tokenizer, model_info))
         except Exception as e:
-            logger.error(f"Failed to load model {model_id}: {str(e)}")
-            yield f"Error loading model {model_id}. Please try again."
-            return
-    
+            logger.error(f"Skipping {model_info.name}: {str(e)}")
+            yield f"⚠️ Couldn't load {model_info.name}, skipping..."
+            continue
+    if not loaded_models:
+        yield "❌ No models could be loaded. Please try again later."
+        return
     responses = []
     formatted_responses = []
     persona_responses = []
-
-    # Generate responses from each council member
-    for i, (persona, (pipe, tokenizer), (model_id, model_name)) in enumerate(zip(selected_personas, loaded_models, selected_models)):
-        display_name = f"{persona['emoji']} {persona['name']} ({model_name})"
-        
-        if i == 0:
-            prompt = create_debate_prompt(user_prompt, persona, debate_style)
-        else:
-            prompt = create_debate_prompt(user_prompt, persona, debate_style, persona_responses)
-        
-        # Stream and collect response
-        response_text = ""
-        for partial in stream_model_response(pipe, tokenizer, prompt, display_name, temperature):
-            # Format the full output
-            current_output = f"**User:** {user_prompt}\n\n" + "\n\n".join(formatted_responses + [partial])
+    if continue_debate and history:
+        formatted_responses.extend(history)
+        persona_responses.extend([r.split("**:")[-1].strip() for r in history if "**:" in r])
+    for i, (persona, (pipe, tokenizer, model_info)) in enumerate(zip(selected_personas, loaded_models)):
+        display_name = f"{persona.emoji} {persona.name} ({model_info.name})"
+        thinking_msg = f"**{display_name}** is thinking..."
+        current_output = "\n\n".join([f"**User:** {user_prompt}"] + formatted_responses + [thinking_msg])
+        yield current_output
+        prompt = create_debate_prompt(
+            user_prompt,
+            persona,
+            debate_style,
+            persona_responses if i > 0 else None
+        )
+        full_response = ""
+        for chunk in stream_response(pipe, tokenizer, prompt, display_name, temperature):
+            full_response = chunk
+            current_output = "\n\n".join([f"**User:** {user_prompt}"] + formatted_responses + [chunk])
             yield current_output
-            response_text = partial.split("**:")[-1].strip()
-        
-        # Add this response to the collected responses
-        persona_responses.append(f"{persona['name']}: {response_text}")
-        formatted_responses.append(partial)
-    
-    # Facilitator synthesis (use a random model)
-    rand_model_idx = random.randint(0, len(loaded_models) - 1)
-    pipe, tokenizer = loaded_models[rand_model_idx]
-    
-    synthesis_prompt = create_synthesis_prompt(user_prompt, persona_responses)
+        persona_responses.append(f"{persona.name}: {full_response.split('**:')[-1].strip()}")
+        formatted_responses.append(full_response)
+    synth_pipe, synth_tokenizer, _ = random.choice(loaded_models)
+    synth_prompt = create_synthesis_prompt(user_prompt, persona_responses)
+    yield "\n\n".join([f"**User:** {user_prompt}"] + formatted_responses + ["✨ **Facilitator** is synthesizing..."])
     synthesis = ""
-    
-    for partial in stream_model_response(pipe, tokenizer, synthesis_prompt, "✨ Facilitator's Synthesis", temperature):
-        current_output = f"**User:** {user_prompt}\n\n" + "\n\n".join(formatted_responses + [partial])
+    for chunk in stream_response(synth_pipe, synth_tokenizer, synth_prompt, "✨ Facilitator", temperature):
+        synthesis = chunk
+        current_output = "\n\n".join([f"**User:** {user_prompt}"] + formatted_responses + [chunk])
         yield current_output
-        synthesis = partial
-    
-    # Final output with timing
     elapsed_time = time.time() - start_time
-    transcript = f"**User:** {user_prompt}\n\n" + "\n\n".join(formatted_responses) + f"\n\n{synthesis}\n\n---\n*Debate completed in {elapsed_time:.1f} seconds*"
+    transcript = (
+        f"**User:** {user_prompt}\n\n" +
+        "\n\n".join(formatted_responses) +
+        f"\n\n{synthesis}\n\n" +
+        f"---\n*Debate completed in {elapsed_time:.1f} seconds*"
+    )
     yield transcript
 
-# Gradio interface with improved UI
+def council_chat_stream_chatbot(
+    user_prompt: str,
+    num_members: int = 3,
+    debate_style: str = "Balanced",
+    temperature: float = 0.7,
+    selected_models: Optional[List[str]] = None,
+    continue_debate: bool = False,
+    history: Optional[List[str]] = None
+) -> Generator[list, None, None]:
+    chat_history = []
+    for output in council_chat_stream(
+        user_prompt, num_members, debate_style, temperature, selected_models, continue_debate, history
+    ):
+        chat_history.append((None, output))
+        yield chat_history
+
 def build_gradio_interface():
-    """Build a more structured and visually appealing Gradio interface"""
-    
-    # Custom CSS for better appearance
     custom_css = """
-    .gradio-container {
-        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-    }
-    .council-header {
-        text-align: center;
-        margin-bottom: 1em;
-    }
-    .council-member {
-        margin: 0.5em 0;
-        padding: 0.5em;
-        border-radius: 8px;
-        background-color: #f5f5f5;
-    }
+    .gradio-container { font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; }
+    #transcript-container { position: relative; }
+    #copy-btn { position: absolute; top: 10px; right: 10px; z-index: 100; }
+    .member-card { border: 1px solid #e0e0e0; border-radius: 8px; padding: 15px; margin-bottom: 15px; background: #f9f9f9; }
+    .member-card h3 { margin-top: 0; color: #333; }
     """
-    
     with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
-        gr.Markdown("# 🤖🏛️ AI Council Debate", elem_classes=["council-header"])
-        gr.Markdown("Ask a question and watch as AI personas debate and deliberate on your topic with different perspectives.")
-        
+        current_debate = gr.State([])
+        gr.Markdown("# 🏛️ AI Council Debate\n*Get diverse AI perspectives on any topic*")
         with gr.Row():
-            with gr.Column():
-                inp = gr.Textbox(
-                    label="Your Topic or Question", 
-                    lines=4, 
-                    placeholder="Enter a topic, question, or issue for the council to debate..."
+            with gr.Column(scale=2):
+                user_prompt = gr.Textbox(
+                    label="Debate Topic",
+                    placeholder="Enter your question or topic for debate...",
+                    lines=4,
+                    max_lines=6
                 )
-                
-                # Advanced options
-                with gr.Accordion("Advanced Options", open=False):
+                with gr.Accordion("⚙️ Debate Settings", open=False):
                     with gr.Row():
                         num_members = gr.Slider(
-                            minimum=2, 
-                            maximum=len(PERSONAS), 
-                            value=3, 
-                            step=1, 
+                            minimum=2,
+                            maximum=len(PERSONAS),
+                            value=3,
+                            step=1,
                             label="Number of Council Members"
                         )
-                        
-                    with gr.Row():
                         debate_style = gr.Radio(
-                            ["Collaborative", "Adversarial", "Balanced"], 
-                            label="Debate Style", 
-                            value="Balanced"
+                            ["Collaborative", "Adversarial", "Balanced"],
+                            value="Balanced",
+                            label="Debate Style"
                         )
-                        
                     with gr.Row():
                         temperature = gr.Slider(
-                            minimum=0.1, 
-                            maximum=1.0, 
-                            value=0.7, 
-                            step=0.1, 
-                            label="Temperature (Creativity)"
+                            minimum=0.1,
+                            maximum=1.0,
+                            value=0.7,
+                            step=0.1,
+                            label="Creativity (Temperature)"
+                        )
+                        model_selection = gr.CheckboxGroup(
+                            choices=[model.name for model in MODELS],
+                            value=[model.name for model in MODELS],
+                            label="Models to Use"
                         )
-                
-                btn = gr.Button("Start Council Debate", variant="primary")
-                
-            with gr.Column():
-                out = gr.Markdown(label="Council Debate Transcript")
-        
-        # Display council members information
-        with gr.Accordion("Meet the Council Members", open=False):
-            member_info = ""
+                    with gr.Row():
+                        continue_btn = gr.Checkbox(
+                            label="Continue Previous Debate",
+                            value=False
+                        )
+                        clear_cache_btn = gr.Button(
+                            "Clear Model Cache",
+                            variant="secondary"
+                        )
+                with gr.Row():
+                    output_style = gr.Radio(
+                        ["Transcript (Markdown)", "Chatbot (Chat History)"],
+                        value="Transcript (Markdown)",
+                        label="Output Style"
+                    )
+                    submit_btn = gr.Button(
+                        "Start Debate",
+                        variant="primary"
+                    )
+                    stop_btn = gr.Button(
+                        "Stop",
+                        variant="stop"
+                    )
+            with gr.Column(scale=3):
+                transcript_out = gr.HTML(label="Council Debate", elem_id="transcript-container", visible=True)
+                chatbot_out = gr.Chatbot(label="Council Debate (Chat)", visible=False, height=500)
+        with gr.Accordion("👥 Meet the Council Members", open=False):
             for persona in PERSONAS:
-                member_info += f"""
-                <div class="council-member">
-                    <h3>{persona['emoji']} {persona['name']}</h3>
-                    <p><strong>Description:</strong> {persona['description']}</p>
-                    <p><strong>Traits:</strong> {persona['traits']}</p>
-                    <p><strong>Communication Style:</strong> {persona['style']}</p>
-                </div>
-                """
-            gr.HTML(member_info)
-        
-        # Example prompts for users to try
-        with gr.Accordion("Example Topics", open=False):
-            examples = [
-                "What role should AI play in education?",
-                "Is universal basic income a good idea?",
-                "How should society balance privacy concerns with security needs?",
-                "What are the ethical implications of genetic engineering?",
-                "How can we address climate change effectively?"
-            ]
-            gr.Examples(examples=examples, inputs=inp)
-        
-        # Event handlers
-        btn.click(
-            fn=council_chat_stream, 
-            inputs=[inp, num_members, debate_style, temperature], 
-            outputs=out
+                with gr.Group(elem_classes="member-card"):
+                    gr.Markdown(f"""
+                    <h3>{persona.emoji} {persona.name}</h3>
+                    <p><strong>Description:</strong> {persona.description}</p>
+                    <p><strong>Traits:</strong> {persona.traits}</p>
+                    <p><strong>Style:</strong> {persona.style}</p>
+                    """)
+        with gr.Accordion("ℹ️ System Information", open=False):
+            gr.Markdown(f"""
+            - **Device:** {'GPU' if torch.cuda.is_available() else 'CPU'}
+            - **Available Models:** {len(MODELS)}
+            - **Council Members:** {len(PERSONAS)}
+            - **Note:** First run may take time to download models
+            """)
+        def route_debate(user_prompt, num_members, debate_style, temperature, model_selection, continue_btn, current_debate, output_style):
+            selected_model_ids = [m.id for m in MODELS if m.name in model_selection]
+            if output_style == "Transcript (Markdown)":
+                for out in council_chat_stream(
+                    user_prompt, num_members, debate_style, temperature, selected_model_ids, continue_btn, current_debate
+                ):
+                    yield gr.update(visible=True, value=out), gr.update(visible=False)
+            else:
+                for out in council_chat_stream_chatbot(
+                    user_prompt, num_members, debate_style, temperature, selected_model_ids, continue_btn, current_debate
+                ):
+                    yield gr.update(visible=False), gr.update(visible=True, value=out)
+        submit_btn.click(
+            route_debate,
+            [user_prompt, num_members, debate_style, temperature, model_selection, continue_btn, current_debate, output_style],
+            [transcript_out, chatbot_out],
+            queue=True
+        )
+        stop_btn.click(
+            fn=None, inputs=None, outputs=None, cancels=[submit_btn]
+        )
+        clear_cache_btn.click(
+            fn=clear_model_cache, inputs=None, outputs=None
+        )
+        def update_history(history: List[str], new_output: str) -> List[str]:
+            if "✨ Facilitator" in new_output:
+                return []
+            return history + [new_output] if history else [new_output]
+        transcript_out.change(
+            fn=update_history,
+            inputs=[current_debate, transcript_out],
+            outputs=current_debate
         )
-        
-        # Footer with additional information
-        gr.Markdown("""
-        ### About This App
-        
-        This application demonstrates how multiple AI models can collaborate in a structured debate.
-        Each AI persona has distinctive traits and perspectives that influence how they approach topics.
-        
-        The models used are open-source LLMs hosted on Hugging Face:
-        - Meta's Llama 3 8B Instruct
-        - Qwen 1.5 7B Chat
-        - Zephyr 7B Beta
-        - Mistral 7B Instruct v0.2
-        
-        ⚠️ Note: First-time loading may take a minute as models are downloaded and initialized.
-        """)
-    
     return demo
 
-# Main application
 if __name__ == "__main__":
-    # Check GPU availability
-    if torch.cuda.is_available():
-        logger.info(f"GPU available: {torch.cuda.get_device_name(0)}")
+    device = get_device()
+    if device == "cuda":
+        gpu_info = torch.cuda.get_device_properties(0)
+        logger.info(f"Using GPU: {gpu_info.name} ({gpu_info.total_memory / (1024**3):.1f}GB)")
     else:
-        logger.info("No GPU available, using CPU. Performance may be slower.")
-    
-    # Create and launch the Gradio interface
+        logger.info("Using CPU")
     demo = build_gradio_interface()
     demo.launch()