Spaces:
Running
Running
| import { pipeline, TextStreamer } from 'https://cdn.jsdelivr.net/npm/@huggingface/[email protected]'; | |
| // DOM Elements | |
| const messagesContainer = document.getElementById('messages'); | |
| const userInput = document.getElementById('userInput'); | |
| const sendBtn = document.getElementById('sendBtn'); | |
| const status = document.getElementById('status'); | |
| const loadingOverlay = document.getElementById('loadingOverlay'); | |
| const loadingText = document.getElementById('loadingText'); | |
| const progressFill = document.getElementById('progressFill'); | |
| const progressText = document.getElementById('progressText'); | |
| // State | |
| let generator = null; | |
| let conversationHistory = [ | |
| { role: "system", content: "You are a helpful, friendly, and knowledgeable AI assistant. Provide clear, concise, and | |
| accurate responses." } | |
| ]; | |
| let isGenerating = false; | |
| // Initialize the application | |
| async function initializeApp() { | |
| try { | |
| updateLoadingStatus('Loading AI model...', 0); | |
| // Create a text generation pipeline with progress callback | |
| generator = await pipeline( | |
| "text-generation", | |
| "onnx-community/Llama-3.2-1B-Instruct-q4f16", | |
| { | |
| dtype: "q4f16", | |
| device: "webgpu", | |
| progress_callback: (progress) => { | |
| if (progress.status === 'progress') { | |
| const percent = Math.round((progress.loaded / progress.total) * 100); | |
| updateLoadingStatus(`Loading ${progress.file}...`, percent); | |
| } else if (progress.status === 'done') { | |
| updateLoadingStatus('Initializing model...', 95); | |
| } | |
| } | |
| } | |
| ); | |
| updateLoadingStatus('Ready!', 100); | |
| // Hide loading overlay after a short delay | |
| setTimeout(() => { | |
| loadingOverlay.classList.add('hidden'); | |
| enableChat(); | |
| }, 500); | |
| } catch (error) { | |
| console.error('Error initializing model:', error); | |
| updateLoadingStatus('Error loading model. Please refresh the page.', 0); | |
| status.textContent = 'Failed to load AI model. Please refresh.'; | |
| status.style.color = '#FF3B30'; | |
| } | |
| } | |
| function updateLoadingStatus(text, percent) { | |
| loadingText.textContent = text; | |
| progressFill.style.width = `${percent}%`; | |
| progressText.textContent = `${percent}%`; | |
| } | |
| function enableChat() { | |
| userInput.disabled = false; | |
| sendBtn.disabled = false; | |
| status.textContent = 'Ready to chat'; | |
| userInput.focus(); | |
| } | |
| function addMessage(role, content) { | |
| // Remove welcome message if it exists | |
| const welcomeMessage = messagesContainer.querySelector('.welcome-message'); | |
| if (welcomeMessage) { | |
| welcomeMessage.remove(); | |
| } | |
| const messageDiv = document.createElement('div'); | |
| messageDiv.className = `message ${role}`; | |
| const contentDiv = document.createElement('div'); | |
| contentDiv.className = 'message-content'; | |
| contentDiv.textContent = content; | |
| messageDiv.appendChild(contentDiv); | |
| messagesContainer.appendChild(messageDiv); | |
| // Scroll to bottom | |
| messagesContainer.scrollTop = messagesContainer.scrollHeight; | |
| return contentDiv; | |
| } | |
| function addTypingIndicator() { | |
| const messageDiv = document.createElement('div'); | |
| messageDiv.className = 'message assistant'; | |
| messageDiv.id = 'typing-indicator'; | |
| const contentDiv = document.createElement('div'); | |
| contentDiv.className = 'message-content'; | |
| const typingDiv = document.createElement('div'); | |
| typingDiv.className = 'typing-indicator'; | |
| typingDiv.innerHTML = ` | |
| <div class="typing-dot"></div> | |
| <div class="typing-dot"></div> | |
| <div class="typing-dot"></div> | |
| `; | |
| contentDiv.appendChild(typingDiv); | |
| messageDiv.appendChild(contentDiv); | |
| messagesContainer.appendChild(messageDiv); | |
| messagesContainer.scrollTop = messagesContainer.scrollHeight; | |
| return messageDiv; | |
| } | |
| function removeTypingIndicator() { | |
| const indicator = document.getElementById('typing-indicator'); | |
| if (indicator) { | |
| indicator.remove(); | |
| } | |
| } | |
| async function generateResponse(userMessage) { | |
| if (isGenerating) return; | |
| isGenerating = true; | |
| sendBtn.disabled = true; | |
| userInput.disabled = true; | |
| status.textContent = 'Thinking...'; | |
| // Add user message to conversation history | |
| conversationHistory.push({ role: "user", content: userMessage }); | |
| // Show typing indicator | |
| const typingIndicator = addTypingIndicator(); | |
| try { | |
| let assistantMessage = ''; | |
| let messageElement = null; | |
| // Create a custom streamer with callback | |
| const streamer = new TextStreamer(generator.tokenizer, { | |
| skip_prompt: true, | |
| skip_special_tokens: true, | |
| callback_function: (text) => { | |
| // Remove typing indicator on first token | |
| if (!messageElement) { | |
| removeTypingIndicator(); | |
| messageElement = addMessage('assistant', ''); | |
| } | |
| // Append the new text | |
| assistantMessage += text; | |
| messageElement.textContent = assistantMessage; | |
| // Scroll to bottom | |
| messagesContainer.scrollTop = messagesContainer.scrollHeight; | |
| } | |
| }); | |
| // Generate response with streaming | |
| const output = await generator(conversationHistory, { | |
| max_new_tokens: 512, | |
| do_sample: false, | |
| temperature: 0.7, | |
| streamer: streamer, | |
| }); | |
| // Get the final response | |
| const finalResponse = output[0].generated_text.at(-1).content; | |
| // Update conversation history | |
| conversationHistory.push({ role: "assistant", content: finalResponse }); | |
| // Ensure the final message is displayed | |
| if (messageElement) { | |
| messageElement.textContent = finalResponse; | |
| } | |
| status.textContent = 'Ready to chat'; | |
| } catch (error) { | |
| console.error('Error generating response:', error); | |
| removeTypingIndicator(); | |
| addMessage('assistant', 'Sorry, I encountered an error. Please try again.'); | |
| status.textContent = 'Error occurred'; | |
| } finally { | |
| isGenerating = false; | |
| sendBtn.disabled = false; | |
| userInput.disabled = false; | |
| userInput.focus(); | |
| } | |
| } | |
| function handleSend() { | |
| const message = userInput.value.trim(); | |
| if (!message || isGenerating) return; | |
| // Add user message to UI | |
| addMessage('user', message); | |
| // Clear input | |
| userInput.value = ''; | |
| userInput.style.height = 'auto'; | |
| // Generate response | |
| generateResponse(message); | |
| } | |
| // Event Listeners | |
| sendBtn.addEventListener('click', handleSend); | |
| userInput.addEventListener('keydown', (e) => { | |
| if (e.key === 'Enter' && !e.shiftKey) { | |
| e.preventDefault(); | |
| handleSend(); | |
| } | |
| }); | |
| // Auto-resize textarea | |
| userInput.addEventListener('input', () => { | |
| userInput.style.height = 'auto'; | |
| userInput.style.height = userInput.scrollHeight + 'px'; | |
| }); | |
| // Initialize the app | |
| initializeApp(); |