Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import torch
+from threading import Thread
 
-# Global variables to store current model and tokenizer
+# Global variables
 current_model = None
 current_tokenizer = None
 current_model_name = None
@@ -12,114 +13,175 @@ def load_model(model_name):
     global current_model, current_tokenizer, current_model_name
     
     if current_model_name == model_name:
-        return f"Model {model_name} is already loaded."
+        return f"Model {model_name.split('/')[-1]} already loaded"
     
     try:
         print(f"Loading model: {model_name}")
         current_tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Optimized loading for speed
         current_model = AutoModelForCausalLM.from_pretrained(
             model_name,
             torch_dtype=torch.float16,
             device_map="auto",
-            low_cpu_mem_usage=True
+            low_cpu_mem_usage=True,
+            use_cache=True  # Enable KV cache for faster generation
         )
+        current_model.eval()  # Set to evaluation mode
         current_model_name = model_name
-        return f"✓ Successfully loaded {model_name}"
+        return f"Model {model_name.split('/')[-1]} loaded successfully"
     except Exception as e:
-        return f"Error loading model: {str(e)}"
+        return f"Error: {str(e)}"
 
-def chat(message, history, model_choice):
-    """Process chat messages"""
+def chat_stream(message, history, model_choice):
+    """Stream chat responses for real-time display"""
     global current_model, current_tokenizer, current_model_name
     
-    # Load model if not loaded or if different model selected
+    # Load model if needed
     if current_model_name != model_choice:
+        yield history + [[message, "Loading model, please wait..."]]
         status = load_model(model_choice)
         if "Error" in status:
-            return status
+            yield history + [[message, status]]
+            return
     
-    # Prepare conversation history
+    # Build conversation history (keep it short for speed)
     conversation = []
-    for user_msg, assistant_msg in history:
+    # Only use last 3 exchanges to keep context short
+    recent_history = history[-3:] if len(history) > 3 else history
+    
+    for user_msg, assistant_msg in recent_history:
         conversation.append({"role": "user", "content": user_msg})
         if assistant_msg:
             conversation.append({"role": "assistant", "content": assistant_msg})
     conversation.append({"role": "user", "content": message})
     
-    # Generate response
     try:
+        # Prepare inputs
         inputs = current_tokenizer.apply_chat_template(
             conversation,
             return_tensors="pt",
             add_generation_prompt=True
         ).to(current_model.device)
         
-        outputs = current_model.generate(
-            inputs,
-            max_new_tokens=512,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=current_tokenizer.eos_token_id
+        # Setup streaming
+        streamer = TextIteratorStreamer(
+            current_tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True
         )
         
-        response = current_tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-        return response
+        # Optimized generation parameters for SPEED
+        generation_kwargs = {
+            "input_ids": inputs,
+            "streamer": streamer,
+            "max_new_tokens": 256,  # Reduced for faster responses
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "top_k": 50,  # Limit vocabulary for speed
+            "do_sample": True,
+            "pad_token_id": current_tokenizer.eos_token_id,
+            "repetition_penalty": 1.1,
+            "num_beams": 1  # No beam search for speed
+        }
+        
+        # Start generation in separate thread
+        thread = Thread(target=current_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        # Stream the response
+        partial_response = ""
+        for new_text in streamer:
+            partial_response += new_text
+            yield history + [[message, partial_response]]
+        
     except Exception as e:
-        return f"Error generating response: {str(e)}"
+        yield history + [[message, f"Error: {str(e)}"]]
 
 # Create Gradio interface
 with gr.Blocks(title="SauLM Legal Assistant", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # ⚖️ SauLM Legal Assistant
-    Chat with specialized legal language models. Select a model and start your conversation.
-    """)
+    
+    gr.Markdown("# SauLM Legal Assistant")
+    gr.Markdown("Ask your legal questions in French or English")
     
     with gr.Row():
-        model_dropdown = gr.Dropdown(
-            choices=[
-                "Equall/SaulLM-54B-Instruct",
-                "Equall/SaulLM-141B-Instruct",
-                "Equall/Saul-7B-Instruct-v1"
-            ],
-            value="Equall/Saul-7B-Instruct-v1",
-            label="Select Model",
-            info="Choose which SauLM model to use"
-        )
+        with gr.Column(scale=3):
+            model_dropdown = gr.Dropdown(
+                choices=[
+                    "Equall/Saul-7B-Instruct-v1",
+                    "Equall/SaulLM-54B-Instruct",
+                    "Equall/SaulLM-141B-Instruct"
+                ],
+                value="Equall/Saul-7B-Instruct-v1",
+                label="Model",
+                info="Start with 7B model for faster responses"
+            )
+        with gr.Column(scale=1):
+            clear = gr.Button("Clear Chat", size="sm")
     
     chatbot = gr.Chatbot(
-        label="Legal Assistant Chat",
-        height=500,
-        show_copy_button=True
+        label="Chat",
+        height=550,
+        show_copy_button=True,
+        bubble_full_width=False
     )
     
     with gr.Row():
         msg = gr.Textbox(
-            label="Your message",
-            placeholder="Ask a legal question...",
-            scale=4
+            label="",
+            placeholder="Type your message here...",
+            scale=9,
+            show_label=False,
+            container=False
         )
         submit = gr.Button("Send", scale=1, variant="primary")
     
-    clear = gr.Button("Clear Chat")
-    
-    gr.Markdown("""
-    ### About SauLM Models
-    - **Saul-7B**: Smallest and fastest model, good for quick responses
-    - **SaulLM-54B**: Balanced performance and quality
-    - **SaulLM-141B**: Most capable model, best for complex legal queries
-    
-    *Note: Larger models require more memory and take longer to load.*
-    """)
+    with gr.Accordion("Model Information", open=False):
+        gr.Markdown("""
+        **Saul-7B-Instruct-v1**
+        - Size: 7 billion parameters
+        - Speed: Fast
+        - Quality: Good
+        - Recommended for quick responses
+        
+        **SaulLM-54B-Instruct**
+        - Size: 54 billion parameters
+        - Speed: Medium
+        - Quality: Very good
+        - Requires more GPU memory
+        
+        **SaulLM-141B-Instruct**
+        - Size: 141 billion parameters
+        - Speed: Slow
+        - Quality: Excellent
+        - Requires significant GPU resources
+        
+        **Note**: Start with the 7B model for faster responses. Larger models provide better quality but are much slower.
+        """)
     
-    # Event handlers
-    msg.submit(chat, [msg, chatbot, model_dropdown], chatbot).then(
-        lambda: "", None, msg
+    # Event handlers with streaming
+    msg.submit(
+        chat_stream,
+        [msg, chatbot, model_dropdown],
+        chatbot
+    ).then(
+        lambda: gr.update(value=""),
+        None,
+        msg
     )
-    submit.click(chat, [msg, chatbot, model_dropdown], chatbot).then(
-        lambda: "", None, msg
+    
+    submit.click(
+        chat_stream,
+        [msg, chatbot, model_dropdown],
+        chatbot
+    ).then(
+        lambda: gr.update(value=""),
+        None,
+        msg
     )
-    clear.click(lambda: None, None, chatbot)
+    
+    clear.click(lambda: [], None, chatbot)
 
 if __name__ == "__main__":
+    demo.queue(max_size=20)  # Queue for streaming
     demo.launch()