Update app.py

app.py

@@ -5,20 +5,33 @@ import threading
 import torch
 import os
 import time
-from typing import List, Dict, Generator, Tuple, Optional
+import sys
 import logging
+from typing import List, Dict, Generator, Tuple, Optional
 from collections import defaultdict
+import gc
+
+# Configure Torch for CPU optimization
+torch.set_num_threads(os.cpu_count() or 1)
+torch.backends.quantized.engine = 'qnnpack' if torch.backends.quantized.supported_engines else None
 
 # Set up logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler('council_debate.log'),
+        logging.StreamHandler()
+    ]
+)
 logger = logging.getLogger(__name__)
 
 # --- Best Free Models for Council ---
 MODELS = [
-    ("mistralai/Mistral-7B-Instruct-v0.2", "Mistral 7B Instruct"),  # Good default choice
-    ("HuggingFaceH4/zephyr-7b-beta", "Zephyr 7B Beta"),  # Smaller alternative
-    ("NousResearch/Hermes-2-Pro-Mistral-7B", "Hermes 2 Pro"),  # Good for debate
-    ("cognitivecomputations/dolphin-2.6-mistral-7b", "Dolphin Mistral"),  # Uncensored
+    ("mistralai/Mistral-7B-Instruct-v0.2", "Mistral 7B Instruct"),
+    ("HuggingFaceH4/zephyr-7b-beta", "Zephyr 7B Beta"),
+    ("NousResearch/Hermes-2-Pro-Mistral-7B", "Hermes 2 Pro"),
+    ("cognitivecomputations/dolphin-2.6-mistral-7b", "Dolphin Mistral"),
 ]
 
 # Define council member personas
@@ -29,7 +42,7 @@ PERSONAS = [
         "traits": "analytical, skeptical, evidence-focused",
         "style": "formal, precise, methodical",
         "emoji": "🔬",
-        "preferred_models": ["Mistral 7B Instruct", "Zephyr 7B Beta"]  # More factual models
+        "preferred_models": ["Mistral 7B Instruct", "Zephyr 7B Beta"]
     },
     {
         "name": "Professor Marcus Chen",
@@ -37,7 +50,7 @@ PERSONAS = [
         "traits": "philosophical, visionary, empathetic",
         "style": "eloquent, metaphorical, conceptual",
         "emoji": "🧠",
-        "preferred_models": ["Hermes 2 Pro", "Dolphin Mistral"]  # More creative models
+        "preferred_models": ["Hermes 2 Pro", "Dolphin Mistral"]
     },
     {
         "name": "Sarah Johnson",
@@ -57,10 +70,9 @@ PERSONAS = [
     }
 ]
 
-# Cache for models to avoid reloading
+# Cache for models
 model_cache = {}
 model_loading_lock = threading.Lock()
-active_sessions = defaultdict(dict)
 stop_signal = threading.Event()
 
 def get_device_preference():
@@ -93,13 +105,15 @@ def load_model(model_id: str) -> Tuple[pipeline, AutoTokenizer]:
                 "torch_dtype": torch.float16 if device == "cuda" else torch.float32
             }
             
-            # More efficient loading for low-memory systems
             if device == "cpu":
-                model_kwargs["low_cpu_mem_usage"] = True
+                model_kwargs.update({
+                    "low_cpu_mem_usage": True,
+                    "torch_dtype": torch.float32,
+                })
             
             model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
             
-            if device != "cuda":  # For CPU/MPS, manually move to device
+            if device != "cuda":
                 model = model.to(device)
             
             pipe = pipeline(
@@ -115,7 +129,6 @@ def load_model(model_id: str) -> Tuple[pipeline, AutoTokenizer]:
         
         except Exception as e:
             logger.error(f"Failed to load model {model_id}: {str(e)}")
-            # Try with smaller precision if failed
             if "out of memory" in str(e).lower() and device == "cuda":
                 logger.info("Attempting to load with float16 to save memory")
                 try:
@@ -128,10 +141,7 @@ def load_model(model_id: str) -> Tuple[pipeline, AutoTokenizer]:
                     logger.error(f"Still failed to load model: {str(e2)}")
             raise
 
-def create_debate_prompt(user_prompt: str, 
-                        persona: Dict, 
-                        debate_style: str = "Balanced", 
-                        previous_responses: Optional[List[str]] = None) -> str:
+def create_debate_prompt(user_prompt: str, persona: Dict, debate_style: str = "Balanced", previous_responses: Optional[List[Dict]] = None) -> str:
     """Enhanced prompt engineering for better debates"""
     persona_desc = (
         f"Roleplay as {persona['name']}, {persona['description']}\n"
@@ -187,12 +197,7 @@ Write in clear, concise bullet points followed by a short paragraph summary.
 
 Facilitator:"""
 
-def stream_model_response(pipe: pipeline, 
-                        tokenizer: AutoTokenizer,
-                        prompt: str, 
-                        speaker_name: str = None,
-                        temperature: float = 0.7,
-                        max_tokens: int = 512) -> Generator[str, None, None]:
+def stream_model_response(pipe: pipeline, tokenizer: AutoTokenizer, prompt: str, speaker_name: str = None, temperature: float = 0.7, max_tokens: int = 512) -> Generator[str, None, None]:
     """Robust streaming with better formatting and stop handling"""
     try:
         if stop_signal.is_set():
@@ -207,7 +212,7 @@ def stream_model_response(pipe: pipeline,
             streamer=streamer,
             max_new_tokens=max_tokens,
             do_sample=True,
-            temperature=min(max(temperature, 0.1), 1.0),  # Clamped to reasonable range
+            temperature=min(max(temperature, 0.1), 1.0),
             top_p=0.95,
             repetition_penalty=1.1,
             eos_token_id=tokenizer.eos_token_id,
@@ -219,19 +224,17 @@ def stream_model_response(pipe: pipeline,
         buffer = ""
         for new_text in streamer:
             if stop_signal.is_set():
-                pipe.model.config.use_cache = False  # Try to stop generation
+                pipe.model.config.use_cache = False
                 thread.join(timeout=1)
                 break
                 
             buffer += new_text
-            # Only yield when we have a complete word to avoid mid-word breaks
             if " " in new_text or "\n" in new_text:
                 if speaker_name:
                     yield f"**{speaker_name}:** {buffer.strip()}"
                 else:
                     yield buffer.strip()
         
-        # Yield any remaining content
         if buffer.strip():
             if speaker_name:
                 yield f"**{speaker_name}:** {buffer.strip()}"
@@ -243,6 +246,9 @@ def stream_model_response(pipe: pipeline,
     except Exception as e:
         logger.error(f"Error in streaming: {str(e)}")
         yield "[Error in generation]" if not speaker_name else f"**{speaker_name}:** [Error in generation]"
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 
 def select_models_for_personas(personas: List[Dict], models: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
     """Match models to personas based on preferences"""
@@ -250,22 +256,16 @@ def select_models_for_personas(personas: List[Dict], models: List[Tuple[str, str
     model_names = [m[1] for m in models]
     
     for persona in personas:
-        # Try to match preferred models first
         for pref in persona.get("preferred_models", []):
             if pref in model_names:
                 selected.append(models[model_names.index(pref)])
                 break
         else:
-            # Fallback to random selection
             selected.append(random.choice(models))
     
     return selected
 
-def council_chat_stream(user_prompt: str, 
-                       num_members: int = 3, 
-                       debate_style: str = "Balanced",
-                       temperature: float = 0.7,
-                       session_id: str = None) -> Generator[str, None, None]:
+def council_chat_stream(user_prompt: str, num_members: int = 3, debate_style: str = "Balanced", temperature: float = 0.7) -> Generator[str, None, None]:
     """Enhanced debate generation with better state management"""
     stop_signal.clear()
     
@@ -276,11 +276,9 @@ def council_chat_stream(user_prompt: str,
     start_time = time.time()
     
     try:
-        # Select personas and models
         selected_personas = random.sample(PERSONAS, min(num_members, len(PERSONAS)))
         selected_models = select_models_for_personas(selected_personas, MODELS)
         
-        # Load all models first with progress updates
         loaded_models = []
         for i, (model_id, model_name) in enumerate(selected_models):
             if stop_signal.is_set():
@@ -300,7 +298,6 @@ def council_chat_stream(user_prompt: str,
             yield "❌ Error: No models could be loaded. Please try again later."
             return
         
-        # Conduct the debate
         responses = []
         formatted_responses = []
         persona_responses = []
@@ -311,17 +308,10 @@ def council_chat_stream(user_prompt: str,
                 return
                 
             display_name = f"{persona['emoji']} {persona['name']} ({model_name})"
-            prompt = create_debate_prompt(            user_prompt, persona, debate_style, persona_responses)
+            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
-            ):
+            for partial in stream_model_response(pipe, tokenizer, prompt, display_name, temperature):
                 if stop_signal.is_set():
                     break
                 yield partial
@@ -331,7 +321,6 @@ def council_chat_stream(user_prompt: str,
                 yield "[Debate stopped during responses]"
                 return
             
-            # Store response data
             response_data = {
                 "name": persona['name'],
                 "model": model_name,
@@ -341,7 +330,6 @@ def council_chat_stream(user_prompt: str,
             persona_responses.append(response_data)
             formatted_responses.append(partial)
         
-        # Facilitator synthesis
         if not stop_signal.is_set():
             yield "\n\n**✨ Council is now synthesizing the discussion...**\n"
             synthesis_model = random.choice(loaded_models)
@@ -352,13 +340,12 @@ def council_chat_stream(user_prompt: str,
                 synthesis_model[1], 
                 synthesis_prompt, 
                 "✨ Facilitator's Synthesis", 
-                temperature*0.8  # Slightly lower temp for synthesis
+                temperature*0.8
             ):
                 if stop_signal.is_set():
                     break
                 yield partial
             
-        # Final output
         elapsed_time = time.time() - start_time
         if not stop_signal.is_set():
             transcript = (
@@ -414,7 +401,6 @@ def build_gradio_interface():
     """
 
     with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
-        # Header section
         with gr.Row():
             gr.Markdown("""
             <div class="council-header">
@@ -423,10 +409,8 @@ def build_gradio_interface():
             </div>
             """)
         
-        # Main controls
         with gr.Row():
             with gr.Column(scale=2):
-                # Input section
                 inp = gr.Textbox(
                     label="Debate Topic",
                     placeholder="Enter a topic or question for the council to debate...",
@@ -434,7 +418,6 @@ def build_gradio_interface():
                     max_lines=6
                 )
                 
-                # Debate controls
                 with gr.Group(elem_classes="debate-controls"):
                     with gr.Row():
                         btn = gr.Button("Start Debate", variant="primary")
@@ -447,8 +430,7 @@ def build_gradio_interface():
                                 minimum=2, 
                                 maximum=4, 
                                 step=1, 
-                                value=3,
-                                info="Number of AI participants"
+                                value=3
                             )
                             debate_style = gr.Dropdown(
                                 label="Debate Style",
@@ -461,11 +443,9 @@ def build_gradio_interface():
                                 minimum=0.1,
                                 maximum=1.0,
                                 step=0.1,
-                                value=0.7,
-                                info="Higher = more creative/random"
+                                value=0.7
                             )
                 
-                # Persona information
                 with gr.Accordion("Meet the Council Members", open=False):
                     for persona in PERSONAS:
                         with gr.Group(elem_classes="persona-card"):
@@ -476,7 +456,6 @@ def build_gradio_interface():
                             **Preferred Models:** {', '.join(persona.get('preferred_models', ['Any']))}
                             """)
             
-            # Output section
             with gr.Column(scale=3):
                 out = gr.Markdown(
                     label="Live Debate Transcript",
@@ -491,7 +470,6 @@ def build_gradio_interface():
                     - Debate memory and context tracking
                     """)
         
-        # Example prompts
         with gr.Accordion("Example Debate Topics", open=False):
             examples = gr.Examples(
                 examples=[
@@ -505,7 +483,6 @@ def build_gradio_interface():
                 label="Click to try these examples"
             )
         
-        # Event handlers
         btn.click(
             fn=council_chat_stream,
             inputs=[inp, num_members, debate_style, temperature],
@@ -517,7 +494,6 @@ def build_gradio_interface():
             queue=False
         )
         
-        # Footer
         gr.Markdown("""
         ---
         **About This System:**
@@ -529,19 +505,36 @@ def build_gradio_interface():
     
     return demo
 
-# Main application
 if __name__ == "__main__":
-    # Check system resources
+    # System checks
     device = get_device_preference()
-    logger.info(f"Running on device: {device}")
+    print(f"\n{'='*40}")
+    print(f"Starting AI Council Debate on {device.upper()}")
+    print(f"Python: {sys.version.split()[0]}")
+    print(f"PyTorch: {torch.__version__}")
+    print(f"Gradio: {gr.__version__}")
+    print(f"{'='*40}\n")
     
     if device == "cpu":
-        logger.warning("Running on CPU - performance will be significantly slower than GPU")
+        print("WARNING: Running on CPU - expect slower performance")
+        print("Recommendations:")
+        print("- Close other memory-intensive applications")
+        print("- Reduce number of council members (2-3)")
+        print("- Be patient with response times (30-90 sec per response)\n")
     
-    # Launch interface
-    demo = build_gradio_interface()
-    demo.queue(concurrency_count=1).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )
+    try:
+        demo = build_gradio_interface()
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False,
+            show_error=True
+        )
+    except Exception as e:
+        print(f"\nERROR: {str(e)}")
+        print("\nTroubleshooting steps:")
+        print("1. Check internet connection (required for model download)")
+        print("2. Verify Hugging Face token is set if using Llama models")
+        print("3. Try reducing number of council members")
+        print("4. Restart the application\n")
+        raise