Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 import torch
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 
 model = None
 tokenizer = None
@@ -15,7 +16,7 @@ alpaca_prompt = """පහත දැක්වෙන්නේ යම් කාර
 {}"""
 
 @spaces.GPU
-def infer(message, history, enable_history=False, max_new_tokens=512):
+def infer(message, history, enable_history=False, max_new_tokens=1024):
     global model, tokenizer
     
     if model is None:
@@ -45,18 +46,52 @@ def infer(message, history, enable_history=False, max_new_tokens=512):
     
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
-    with torch.inference_mode():
-        outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
+    # Initialize the streamer
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Generation parameters
+    generation_kwargs = dict(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        streamer=streamer,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.95,
+    )
+    
+    # Start generation in a separate thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
     
-    if "### ප්‍රතිචාරය:" in text:
-        text = text.split("### ප්‍රතිචාරය:")[-1].strip()
+    # Stream the response
+    partial_text = ""
+    response_started = False
     
-    return text
+    for new_text in streamer:
+        partial_text += new_text
+        
+        # Check if we've reached the response section
+        if "### ප්‍රතිචාරය:" in partial_text and not response_started:
+            partial_text = partial_text.split("### ප්‍රතිචාරය:")[-1].strip()
+            response_started = True
+        
+        if response_started:
+            yield partial_text
 
-# Custom CSS for styling
+# Custom CSS for styling with copy button
 custom_css = """
+/* Container width constraints for PC screens */
+.gradio-container {
+    max-width: 1200px !important;
+    margin: 0 auto !important;
+}
+
+/* Chat interface max width */
+.contain {
+    max-width: 900px !important;
+    margin: 0 auto !important;
+}
+
 #splash-screen {
     position: fixed;
     top: 0;
@@ -93,6 +128,42 @@ custom_css = """
     75% { transform: rotateY(540deg) scale(1.2); opacity: 0.8; }
 }
 
+/* Copy button styling */
+.message-wrap.bot {
+    position: relative;
+}
+
+.message-wrap.bot:hover .copy-button {
+    opacity: 1;
+}
+
+.copy-button {
+    position: absolute;
+    top: 8px;
+    right: 8px;
+    opacity: 0;
+    transition: opacity 0.2s;
+    background: rgba(102, 126, 234, 0.1);
+    border: 1px solid rgba(102, 126, 234, 0.3);
+    border-radius: 6px;
+    padding: 6px 10px;
+    cursor: pointer;
+    font-size: 12px;
+    color: #667eea;
+    z-index: 10;
+}
+
+.copy-button:hover {
+    background: rgba(102, 126, 234, 0.2);
+    border-color: rgba(102, 126, 234, 0.5);
+}
+
+.copy-button.copied {
+    background: rgba(34, 197, 94, 0.2);
+    border-color: rgba(34, 197, 94, 0.5);
+    color: #22c55e;
+}
+
 /* Smaller font sizes for chat */
 .message-wrap .message {
     font-size: 0.9rem !important;
@@ -102,6 +173,15 @@ custom_css = """
     font-size: 0.9rem !important;
 }
 
+/* Compact examples grid */
+.examples {
+    max-width: 100% !important;
+}
+
+.examples .wrap {
+    gap: 0.5rem !important;
+}
+
 /* Avatar styling */
 .message-wrap.user .avatar-container::before {
     content: "👤";
@@ -137,6 +217,67 @@ custom_css = """
     height: 40px !important;
     min-width: 40px !important;
 }
+
+/* Compact padding */
+.main {
+    padding: 1rem !important;
+}
+
+/* Title styling - more compact */
+h1 {
+    font-size: 1.5rem !important;
+    margin-bottom: 0.5rem !important;
+}
+
+/* Streaming cursor effect */
+.message-wrap.bot.streaming .message::after {
+    content: '▊';
+    animation: blink 1s infinite;
+    margin-left: 2px;
+}
+
+@keyframes blink {
+    0%, 50% { opacity: 1; }
+    51%, 100% { opacity: 0; }
+}
+"""
+
+# JavaScript for copy functionality
+copy_js = """
+<script>
+function addCopyButtons() {
+    // Remove existing copy buttons first
+    document.querySelectorAll('.copy-button').forEach(btn => btn.remove());
+    
+    // Add copy buttons to bot messages
+    const botMessages = document.querySelectorAll('.message-wrap.bot .message');
+    botMessages.forEach((message, index) => {
+        if (!message.querySelector('.copy-button')) {
+            const copyBtn = document.createElement('button');
+            copyBtn.className = 'copy-button';
+            copyBtn.innerHTML = '📋 Copy';
+            copyBtn.onclick = function(e) {
+                e.stopPropagation();
+                const text = message.innerText;
+                navigator.clipboard.writeText(text).then(() => {
+                    copyBtn.innerHTML = '✅ Copied!';
+                    copyBtn.classList.add('copied');
+                    setTimeout(() => {
+                        copyBtn.innerHTML = '📋 Copy';
+                        copyBtn.classList.remove('copied');
+                    }, 2000);
+                });
+            };
+            message.parentElement.style.position = 'relative';
+            message.parentElement.appendChild(copyBtn);
+        }
+    });
+}
+
+// Run on load and periodically to catch new messages
+setInterval(addCopyButtons, 1000);
+window.addEventListener('load', addCopyButtons);
+</script>
 """
 
 # Splash screen HTML
@@ -152,12 +293,12 @@ splash_html = """
 """
 
 # ---------------- UI ----------------
-with gr.Blocks(css=custom_css) as demo:
+with gr.Blocks(css=custom_css, head=copy_js) as demo:
     gr.HTML(splash_html)
     
     chat = gr.ChatInterface(
         fn=lambda message, history: infer(message, history, enable_history.value, max_new_tokens.value),
-        title="🧠 Gamunu 4B Instruct",
+        title="🧠 Gamunu 4B Instruct - Demo",
         theme=gr.themes.Default(text_size="sm"),
         examples=[
             ["හෙලෝ ගැමුණු! මම සමන්, ඔයාට කොහොමද?"],