Update index.js

index.js

@@ -1,11 +1,12 @@
 import { pipeline, TextStreamer } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.3';
 
-class Chatbot {
+class StreamingChatbot {
     constructor() {
         this.generator = null;
         this.messages = [];
         this.isProcessing = false;
         this.useWebGPU = false;
+        this.currentStreamer = null;
         this.initElements();
         this.initEventListeners();
         this.initModel();
@@ -50,7 +51,6 @@ class Chatbot {
             }
         });
 
-        // Check WebGPU support
         this.checkWebGPUSupport();
     }
 
@@ -77,6 +77,9 @@ class Chatbot {
     updateCharCount() {
         const count = this.userInput.value.length;
         this.charCount.textContent = `${count} / 1000`;
+        if (this.charCount) {
+            this.charCount.style.color = count > 900 ? '#ef4444' : '#64748b';
+        }
     }
 
     autoResize() {
@@ -141,60 +144,77 @@ class Chatbot {
             const responseContainer = this.addMessage('', 'assistant', true);
             const messageContent = responseContainer.querySelector('.message-content');
             
-            // Custom streamer to update UI
-            const streamer = {
-                tokenizer: this.generator.tokenizer,
+            // Add streaming indicator
+            const streamingIndicator = document.createElement('span');
+            streamingIndicator.className = 'streaming-indicator';
+            streamingIndicator.textContent = '●';
+            streamingIndicator.style.animation = 'pulse 1.5s infinite';
+            streamingIndicator.style.color = '#10b981';
+            streamingIndicator.style.marginLeft = '4px';
+            messageContent.appendChild(streamingIndicator);
+
+            // Create enhanced TextStreamer with proper callback handling
+            this.currentStreamer = new TextStreamer(this.generator.tokenizer, {
                 skip_prompt: true,
                 skip_special_tokens: true,
-                text_cache: '',
-                put: function(value) {
-                    const decoded = this.tokenizer.decode(value[0], { 
-                        skip_special_tokens: this.skip_special_tokens 
-                    });
-                    if (this.skip_prompt && !this.started) {
-                        this.started = true;
-                        return;
-                    }
-                    this.text_cache += decoded;
-                    // Only update when we have complete words (space or punctuation)
-                    if (/[.!?;:\s\n]$/.test(this.text_cache)) {
-                        messageContent.textContent += this.text_cache;
-                        this.scrollToBottom();
-                        this.text_cache = '';
+                callback_function: (text) => {
+                    // Remove streaming indicator on first text
+                    if (streamingIndicator.parentNode) {
+                        streamingIndicator.remove();
                     }
-                },
-                end: function() {
-                    if (this.text_cache) {
-                        messageContent.textContent += this.text_cache;
-                        this.scrollToBottom();
-                        this.text_cache = '';
-                    }
-                },
-                on_finalized_text: function(text, stream_end) {
+                    // Append new text and scroll
                     messageContent.textContent += text;
                     this.scrollToBottom();
-                },
-                scrollToBottom: () => this.scrollToBottom()
-            };
+                }
+            });
+
+            // Prepare conversation context
+            const conversationText = this.messages.map(msg => 
+                `${msg.role === 'user' ? 'User' : 'Assistant'}: ${msg.content}`
+            ).join('\n') + '\nAssistant:';
 
-            // Generate response
-            const output = await this.generator(this.messages, {
-                max_new_tokens: 500,
-                do_sample: false,
-                streamer: streamer,
+            // Generate response with streaming
+            const output = await this.generator(conversationText, {
+                max_new_tokens: 300,
+                do_sample: true,
+                temperature: 0.7,
+                top_p: 0.9,
+                streamer: this.currentStreamer,
+                return_full_text: false
             });
 
-            const assistantMessage = output[0].generated_text.at(-1).content;
-            this.messages.push({ role: "assistant", content: assistantMessage });
+            // Clean up streaming indicator if still present
+            if (streamingIndicator.parentNode) {
+                streamingIndicator.remove();
+            }
+
+            // Get the generated text and add to conversation history
+            const generatedText = output[0].generated_text || messageContent.textContent;
+            if (generatedText && !messageContent.textContent) {
+                messageContent.textContent = generatedText;
+            }
+
+            // Add assistant's response to message history
+            const finalResponse = messageContent.textContent.trim();
+            if (finalResponse) {
+                this.messages.push({ role: "assistant", content: finalResponse });
+            }
             
         } catch (error) {
             console.error('Error generating response:', error);
-            this.addMessage('Sorry, I encountered an error. Please try again.', 'error');
+            // Remove any existing content and show error
+            const responseContainer = this.chatMessages.lastElementChild;
+            if (responseContainer && responseContainer.classList.contains('assistant')) {
+                responseContainer.remove();
+            }
+            this.addMessage('Sorry, I encountered an error while generating the response. Please try again.', 'error');
         } finally {
             this.loadingIndicator.classList.add('hidden');
             this.isProcessing = false;
             this.sendButton.disabled = false;
+            this.currentStreamer = null;
             this.userInput.focus();
+            this.scrollToBottom();
         }
     }
 
@@ -226,32 +246,54 @@ class Chatbot {
     }
 
     scrollToBottom() {
-        this.chatMessages.scrollTop = this.chatMessages.scrollHeight;
+        requestAnimationFrame(() => {
+            this.chatMessages.scrollTop = this.chatMessages.scrollHeight;
+        });
     }
 
     clearChat() {
+        if (this.isProcessing) {
+            this.showNotification('Please wait for the current response to complete');
+            return;
+        }
+        
         this.messages = [];
         this.chatMessages.innerHTML = `
             <div class="welcome-message">
                 <div class="welcome-icon">🤖</div>
                 <h2>Welcome to AI Assistant</h2>
-                <p>I'm powered by MobileLLM, a lightweight language model. Ask me anything!</p>
+                <p>I'm powered by MobileLLM with real-time streaming. Ask me anything!</p>
             </div>
         `;
+        this.showNotification('Chat cleared!', 'info');
     }
 
     showNotification(message, type = 'info') {
         const notification = document.createElement('div');
         notification.className = `notification ${type}`;
         notification.textContent = message;
+        notification.style.cssText = `
+            position: fixed;
+            top: 20px;
+            right: 20px;
+            padding: 12px 20px;
+            border-radius: 8px;
+            color: white;
+            font-weight: 500;
+            transform: translateX(100%);
+            transition: all 0.3s;
+            z-index: 1001;
+            background: ${type === 'success' ? '#10b981' : type === 'error' ? '#ef4444' : '#3b82f6'};
+        `;
+        
         document.body.appendChild(notification);
         
         setTimeout(() => {
-            notification.classList.add('show');
+            notification.style.transform = 'translateX(0)';
         }, 10);
         
         setTimeout(() => {
-            notification.classList.remove('show');
+            notification.style.transform = 'translateX(100%)';
             setTimeout(() => notification.remove(), 300);
         }, 3000);
     }
@@ -259,5 +301,5 @@ class Chatbot {
 
 // Initialize chatbot when DOM is ready
 document.addEventListener('DOMContentLoaded', () => {
-    new Chatbot();
+    new StreamingChatbot();
 });