Update app.py

app.py

@@ -8,20 +8,36 @@ import time
 from typing import List, Dict, Generator, Tuple, Optional, Union
 import logging
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 import gc
+from enum import Enum
+import json
+from pathlib import Path
+import uuid
+from datetime import datetime
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
-
 warnings.filterwarnings("ignore", message="torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly")
 
+# Enums and Data Classes
+class DebateStyle(str, Enum):
+    COLLABORATIVE = "Collaborative"
+    ADVERSARIAL = "Adversarial"
+    BALANCED = "Balanced"
+
+class OutputStyle(str, Enum):
+    TRANSCRIPT = "Transcript (Markdown)"
+    CHATBOT = "Chatbot (Chat History)"
+
 @dataclass
 class ModelInfo:
     id: str
     name: str
     required_memory: str  # Estimated VRAM requirement
+    supports_quantization: bool = False
+    quantization_config: Optional[Dict] = field(default_factory=dict)
 
 @dataclass
 class Persona:
@@ -31,11 +47,103 @@ class Persona:
     style: str
     emoji: str
 
+@dataclass
+class DebateHistoryItem:
+    id: str
+    timestamp: float
+    topic: str
+    transcript: str
+    participants: List[str]
+    debate_style: str
+
+class ModelPerformance:
+    def __init__(self):
+        self.times = {}
+        self.token_counts = {}
+    
+    def record_generation(self, model_id: str, time_taken: float, tokens_generated: int):
+        if model_id not in self.times:
+            self.times[model_id] = []
+            self.token_counts[model_id] = []
+        self.times[model_id].append(time_taken)
+        self.token_counts[model_id].append(tokens_generated)
+    
+    def get_stats(self, model_id: str) -> Dict:
+        times = self.times.get(model_id, [])
+        tokens = self.token_counts.get(model_id, [])
+        return {
+            "avg_time": sum(times) / max(1, len(times)),
+            "avg_tokens": sum(tokens) / max(1, len(tokens)),
+            "total_calls": len(times),
+            "tokens_per_second": sum(tokens) / max(1, sum(times)) if times else 0
+        }
+
+class DebateHistoryManager:
+    HISTORY_FILE = "debate_history.json"
+    
+    @classmethod
+    def save_history(cls, item: DebateHistoryItem):
+        history = cls.load_history()
+        history.append({
+            "id": item.id,
+            "timestamp": item.timestamp,
+            "topic": item.topic,
+            "transcript": item.transcript,
+            "participants": item.participants,
+            "debate_style": item.debate_style
+        })
+        with open(cls.HISTORY_FILE, "w") as f:
+            json.dump(history, f, indent=2)
+    
+    @classmethod
+    def load_history(cls) -> List[Dict]:
+        path = Path(cls.HISTORY_FILE)
+        if not path.exists():
+            return []
+        try:
+            with open(path, "r") as f:
+                return json.load(f)
+        except Exception as e:
+            logger.error(f"Error loading history: {str(e)}")
+            return []
+    
+    @classmethod
+    def get_history_item(cls, item_id: str) -> Optional[Dict]:
+        history = cls.load_history()
+        for item in history:
+            if item["id"] == item_id:
+                return item
+        return None
+
+# Constants
 MODELS = [
-    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"),
+    ModelInfo(
+        "meta-llama/Meta-Llama-3-8B-Instruct",
+        "Llama 3 8B Instruct",
+        "16GB",
+        True,
+        {"load_in_4bit": True, "bnb_4bit_compute_dtype": torch.float16}
+    ),
+    ModelInfo(
+        "Qwen/Qwen1.5-7B-Chat",
+        "Qwen1.5 7B Chat",
+        "14GB",
+        True,
+        {"load_in_4bit": True, "bnb_4bit_compute_dtype": torch.float16}
+    ),
+    ModelInfo(
+        "HuggingFaceH4/zephyr-7b-beta",
+        "Zephyr 7B Beta",
+        "14GB",
+        False
+    ),
+    ModelInfo(
+        "mistralai/Mistral-7B-Instruct-v0.2",
+        "Mistral 7B Instruct",
+        "14GB",
+        True,
+        {"load_in_4bit": True, "bnb_4bit_compute_dtype": torch.float16}
+    ),
 ]
 
 PERSONAS = [
@@ -69,9 +177,12 @@ PERSONAS = [
     )
 ]
 
+# Global State
 model_cache = {}
 current_device = None
+performance_monitor = ModelPerformance()
 
+# Core Functions
 def get_device() -> str:
     global current_device
     if current_device:
@@ -102,32 +213,39 @@ def load_model(model_info: ModelInfo) -> Tuple[pipeline, AutoTokenizer]:
     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()
+    kwargs = {"trust_remote_code": True}
+    
     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:
+        if gpu_mem < required_mem and not model_info.supports_quantization:
             logger.warning(f"Insufficient GPU memory for {model_info.name} (needs {required_mem}GB, has {gpu_mem:.1f}GB)")
+    
+    # Handle quantization if supported and on CUDA
+    if device == "cuda" and model_info.supports_quantization:
+        kwargs.update(model_info.quantization_config)
+        kwargs["device_map"] = "auto"
+    else:
+        kwargs["torch_dtype"] = torch.float16 if device == "cuda" else torch.float32
+    
     logger.info(f"Loading {model_info.name} on {device}")
     try:
         start_time = time.time()
         tokenizer = AutoTokenizer.from_pretrained(model_info.id)
-        model_kwargs = {
-            "trust_remote_code": True,
-            "device_map": "auto" if device == "cuda" else None,
-            "torch_dtype": torch.float16 if device == "cuda" else torch.float32
-        }
-        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 = AutoModelForCausalLM.from_pretrained(model_info.id, **kwargs)
+        
+        if device == "cuda" and not model_info.supports_quantization:
+            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")
@@ -139,14 +257,15 @@ def load_model(model_info: ModelInfo) -> Tuple[pipeline, AutoTokenizer]:
 def create_debate_prompt(
     user_prompt: str,
     persona: Persona,
-    debate_style: str = "Balanced",
+    debate_style: DebateStyle = DebateStyle.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."
+        DebateStyle.COLLABORATIVE: "Focus on building upon ideas and finding common ground.",
+        DebateStyle.ADVERSARIAL: "Challenge assumptions and present strong contrasting views.",
+        DebateStyle.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}
@@ -155,6 +274,7 @@ You're in a council debating: "{user_prompt}"
 
 {style_guidance}
 Respond naturally in 3-4 paragraphs."""
+    
     if previous_responses:
         debate_history = "\n\n".join(previous_responses)
         return f"""{base_prompt}
@@ -164,6 +284,7 @@ Current discussion:
 
 Now respond thoughtfully to the ongoing debate:
 {persona.name}:"""
+    
     return f"""{base_prompt}
 
 Begin your response:
@@ -193,9 +314,13 @@ def stream_response(
     temperature: float = 0.7,
     max_tokens: int = 512
 ) -> Generator[str, None, None]:
+    start_time = time.time()
+    tokens_generated = 0
+    
     try:
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
         input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(pipe.model.device)
+        
         generation_kwargs = dict(
             input_ids=input_ids,
             streamer=streamer,
@@ -206,22 +331,36 @@ def stream_response(
             repetition_penalty=1.1,
             eos_token_id=tokenizer.eos_token_id,
         )
+        
         thread = threading.Thread(target=pipe.model.generate, kwargs=generation_kwargs)
         thread.start()
+        
         buffer = ""
         for new_text in streamer:
             buffer += new_text
+            tokens_generated += len(tokenizer.encode(new_text))
             if new_text and new_text[-1] in " .,;!?\n":
                 if speaker_name:
                     yield f"**{speaker_name}:** {buffer.strip()}"
                 else:
                     yield buffer.strip()
+                buffer = ""
+        
         thread.join()
+        
         if buffer.strip():
             if speaker_name:
                 yield f"**{speaker_name}:** {buffer.strip()}"
             else:
                 yield buffer.strip()
+        
+        # Record performance metrics
+        performance_monitor.record_generation(
+            pipe.model.config._name_or_path,
+            time.time() - start_time,
+            tokens_generated
+        )
+        
     except Exception as e:
         logger.error(f"Streaming error: {str(e)}")
         error_msg = f"[Error: {str(e)}]"
@@ -230,21 +369,32 @@ def stream_response(
 def council_chat_stream(
     user_prompt: str,
     num_members: int = 3,
-    debate_style: str = "Balanced",
+    debate_style: Union[DebateStyle, str] = DebateStyle.BALANCED,
     temperature: float = 0.7,
     selected_models: Optional[List[str]] = None,
     continue_debate: bool = False,
-    history: Optional[List[str]] = None
+    history: Optional[List[str]] = None,
+    save_history: bool = True
 ) -> Generator[str, None, None]:
     if not user_prompt.strip():
         yield "Please enter a topic for debate."
         return
+    
+    # Convert string style to Enum if needed
+    if isinstance(debate_style, str):
+        try:
+            debate_style = DebateStyle(debate_style)
+        except ValueError:
+            debate_style = DebateStyle.BALANCED
+    
     num_members = max(2, min(num_members, len(PERSONAS)))
     temperature = max(0.1, min(temperature, 1.0))
+    
     start_time = time.time()
     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_info in selected_model_infos:
         try:
@@ -256,41 +406,56 @@ def council_chat_stream(
             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 = []
+    participant_names = []
+    
     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})"
+        participant_names.append(display_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
+        
         persona_responses.append(f"{persona.name}: {full_response.split('**:')[-1].strip()}")
         formatted_responses.append(full_response)
+    
+    # Generate synthesis
     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 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
+    
     elapsed_time = time.time() - start_time
     transcript = (
         f"**User:** {user_prompt}\n\n" +
@@ -298,153 +463,119 @@ def council_chat_stream(
         f"\n\n{synthesis}\n\n" +
         f"---\n*Debate completed in {elapsed_time:.1f} seconds*"
     )
+    
+    # Save to history
+    if save_history:
+        history_item = DebateHistoryItem(
+            id=str(uuid.uuid4()),
+            timestamp=time.time(),
+            topic=user_prompt,
+            transcript=transcript,
+            participants=participant_names,
+            debate_style=debate_style.value
+        )
+        DebateHistoryManager.save_history(history_item)
+    
     yield transcript
 
 def council_chat_stream_chatbot(
     user_prompt: str,
     num_members: int = 3,
-    debate_style: str = "Balanced",
+    debate_style: Union[DebateStyle, str] = DebateStyle.BALANCED,
     temperature: float = 0.7,
     selected_models: Optional[List[str]] = None,
     continue_debate: bool = False,
-    history: Optional[List[str]] = None
+    history: Optional[List[str]] = None,
+    save_history: bool = True
 ) -> Generator[list, None, None]:
     chat_history = []
     for output in council_chat_stream(
-        user_prompt, num_members, debate_style, temperature, selected_models, continue_debate, history
+        user_prompt, num_members, debate_style, temperature, 
+        selected_models, continue_debate, history, save_history
     ):
         chat_history.append((None, output))
         yield chat_history
 
+# UI Components
+def build_persona_card(persona: Persona) -> gr.Box:
+    with gr.Box(elem_classes="member-card") as 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>
+        """)
+    return card
+
+def build_model_info_card(model: ModelInfo) -> gr.Box:
+    with gr.Box(elem_classes="model-card") as card:
+        gr.Markdown(f"""
+        <h3>{model.name}</h3>
+        <p><strong>ID:</strong> {model.id}</p>
+        <p><strong>Memory Requirement:</strong> {model.required_memory}</p>
+        <p><strong>Quantization:</strong> {'Supported' if model.supports_quantization else 'Not Supported'}</p>
+        """)
+    return card
+
+def build_history_item_ui(history_item: Dict) -> gr.Box:
+    with gr.Box(elem_classes="history-item") as item:
+        with gr.Row():
+            with gr.Column(scale=3):
+                gr.Markdown(f"**{history_item['topic']}**")
+                gr.Markdown(f"*{datetime.fromtimestamp(history_item['timestamp']).strftime('%Y-%m-%d %H:%M:%S')}*")
+            with gr.Column(scale=1):
+                view_btn = gr.Button("View", size="sm")
+                load_btn = gr.Button("Load", size="sm")
+    return item, view_btn, load_btn
+
+# Gradio Interface
 def build_gradio_interface():
     custom_css = """
     .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; }
+    .member-card, .model-card, .history-item { 
+        border: 1px solid #e0e0e0; 
+        border-radius: 8px; 
+        padding: 15px; 
+        margin-bottom: 15px; 
+        background: #f9f9f9; 
+    }
+    .member-card h3, .model-card h3 { margin-top: 0; color: #333; }
+    #transcript-container { position: relative; max-height: 600px; overflow-y: auto; }
+    #chatbot-container { max-height: 600px; }
+    .stats-table { width: 100%; border-collapse: collapse; }
+    .stats-table th, .stats-table td { padding: 8px; text-align: left; border-bottom: 1px solid #ddd; }
     """
+    
     with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
         current_debate = gr.State([])
+        current_history_id = gr.State(None)
+        
         gr.Markdown("# 🏛️ AI Council Debate\n*Get diverse AI perspectives on any topic*")
+        
         with gr.Row():
             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
-                )
-                with gr.Accordion("⚙️ Debate Settings", open=False):
-                    with gr.Row():
-                        num_members = gr.Slider(
-                            minimum=2,
-                            maximum=len(PERSONAS),
-                            value=3,
-                            step=1,
-                            label="Number of Council Members"
-                        )
-                        debate_style = gr.Radio(
-                            ["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="Creativity (Temperature)"
-                        )
-                        model_selection = gr.CheckboxGroup(
-                            choices=[model.name for model in MODELS],
-                            value=[model.name for model in MODELS],
-                            label="Models to Use"
-                        )
-                    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"
+                # Debate Input Section
+                with gr.Group():
+                    user_prompt = gr.Textbox(
+                        label="Debate Topic",
+                        placeholder="Enter your question or topic for debate...",
+                        lines=4,
+                        max_lines=6
                     )
-                    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:
-                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
-        )
-    return demo
-
-if __name__ == "__main__":
-    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("Using CPU")
-    demo = build_gradio_interface()
-    demo.launch()
+                    
+                    with gr.Accordion("⚙️ Debate Settings", open=False):
+                        with gr.Row():
+                            num_members = gr.Slider(
+                                minimum=2,
+                                maximum=len(PERSONAS),
+                                value=3,
+                                step=1,
+                                label="Number of Council Members"
+                            )
+                            debate_style = gr.Radio(
+                                list(DebateStyle),
+                                value=DebateStyle.BALANCED,
+                                label="Debate Style"
+                            )
+                        
+