show context window usage

Author: letmaik

src/App.jsx

@@ -45,6 +45,12 @@ function setStoredModel(modelId) {
   }
 }
 
+// Get context window size for different models
+function getContextWindowSize(modelUrl) {
+  const model = AVAILABLE_MODELS.find(m => m.url === modelUrl);
+  return model?.contextSize || 8192; // Default to 8k if unknown
+}
+
 function App() {
   // Create a reference to the worker object.
   const worker = useRef(null);
@@ -68,6 +74,7 @@ function App() {
   const [queuedMessage, setQueuedMessage] = useState(null); // For storing message when model is loading
   const [tps, setTps] = useState(null);
   const [numTokens, setNumTokens] = useState(null);
+  const [contextTokens, setContextTokens] = useState(null); // Current context window usage
 
   // Mobile detection
   const [isMobile, setIsMobile] = useState(false);
@@ -249,9 +256,10 @@ function App() {
           {
             // Generation update: update the output text.
             // Parse messages
-            const { output, tps, numTokens } = e.data;
+            const { output, tps, numTokens, contextTokens } = e.data;
             setTps(tps);
             setNumTokens(numTokens);
+            setContextTokens(contextTokens);
             setMessages((prev) => {
               const cloned = [...prev];
               const last = cloned.at(-1);
@@ -489,6 +497,17 @@ function App() {
                 {!isRunning && (
                   <span className="mr-1">&#41;.</span>
                 )}
+                {contextTokens && (
+                  <>
+                    <span className="mx-2">•</span>
+                    <span className="text-gray-500 dark:text-gray-300">
+                      Context: {contextTokens.toLocaleString()}/{getContextWindowSize(selectedModel).toLocaleString()} tokens
+                    </span>
+                    <span className="text-gray-400 dark:text-gray-500 text-xs ml-1">
+                      ({((contextTokens / getContextWindowSize(selectedModel)) * 100).toFixed(1)}%)
+                    </span>
+                  </>
+                )}
               </>
             )}
           </p>

src/components/ModelSelector.jsx

@@ -6,42 +6,48 @@ const AVAILABLE_MODELS = [
     name: 'Llama 3.2 1B',
     description: 'Meta model, 1.2 GB',
     url: 'onnx-community/Llama-3.2-1B-Instruct-q4f16',
-    hasReasoningBlocks: false
+    hasReasoningBlocks: false,
+    contextSize: 131072 // 128k tokens
   },
   {
     id: 'llama-3.2-3b',
     name: 'Llama 3.2 3B',
     description: 'Meta model, 2.4 GB',
     url: 'onnx-community/Llama-3.2-3B-Instruct-onnx-web-gqa',
-    hasReasoningBlocks: false
+    hasReasoningBlocks: false,
+    contextSize: 131072 // 128k tokens
   },
   {
     id: 'phi-3.5-mini',
     name: 'Phi-3.5 Mini 3.8B',
     description: 'Microsoft model, 2.1 GB',
     url: 'onnx-community/Phi-3.5-mini-instruct-onnx-web',
-    hasReasoningBlocks: false
+    hasReasoningBlocks: false,
+    contextSize: 131072 // 128k tokens
   },
   {
     id: 'smollm2-1.7b',
     name: 'SmolLM2 1.7B',
     description: 'HuggingFace model, 1.1 GB',
     url: 'HuggingFaceTB/SmolLM2-1.7B-Instruct',
-    hasReasoningBlocks: false
+    hasReasoningBlocks: false,
+    contextSize: 8192 // 8k tokens
   },
   {
     id: 'qwen3-0.6b',
     name: 'Qwen3 0.6B',
     description: 'Alibaba model, 0.5 GB',
     url: 'onnx-community/Qwen3-0.6B-ONNX',
-    hasReasoningBlocks: true
+    hasReasoningBlocks: true,
+    contextSize: 32768 // 32k tokens
   },
   {
     id: 'deepseek-r1-distill-qwen-1.5b',
     name: 'DeepSeek-R1-Distill-Qwen 1.5B',
     description: 'DeepSeek model, 1.3 GB',
     url: 'onnx-community/DeepSeek-R1-Distill-Qwen-1.5B-ONNX',
-    hasReasoningBlocks: true
+    hasReasoningBlocks: true,
+    contextSize: 131072 // 128k tokens
   }
 ];
 

src/worker.js

@@ -84,6 +84,9 @@ async function generate(messages, model_id) {
     return_dict: true,
   });
 
+  // Calculate input token count for context tracking
+  const inputTokenCount = inputs.input_ids.dims[1]; // Get the sequence length
+
   let startTime;
   let numTokens = 0;
   let tps;
@@ -100,6 +103,7 @@ async function generate(messages, model_id) {
       output,
       tps,
       numTokens,
+      contextTokens: inputTokenCount + numTokens, // Total tokens used in context
     });
   };