Update app.py

app.py

@@ -18,27 +18,20 @@ ONNX_MODEL_FILE = "model.onnx"
 # Shared tokenizer
 tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID, token=token)
 
-def greedy_decode_onnx(session, input_ids, attention_mask, max_new_tokens=50):
-    generated = input_ids.copy()
-    for _ in range(max_new_tokens):
-        outputs = session.run(None, {
-            "input_ids": generated,
-            "attention_mask": attention_mask
-        })
-        next_token_logits = outputs[0][:, -1, :]
-        next_token = np.argmax(next_token_logits, axis=-1).reshape(-1, 1)
-        generated = np.concatenate((generated, next_token), axis=1)
-        attention_mask = np.concatenate(
-            (attention_mask, np.ones((1, 1), dtype=np.int64)), axis=1)
-        if next_token[0][0] == tokenizer.eos_token_id:
-            break
-    return tokenizer.decode(generated[0], skip_special_tokens=True)
-
-def compare_outputs(prompt):
+def compare_outputs(prompt, show_tokens):
     summary_log = []
+    pt_output_text = ""
+    ort_output_text = ""
+    pt_tokens = []
+    ort_tokens = []
+
+    try:
+        import psutil
+        ram_used = f"{psutil.virtual_memory().used / 1e9:.2f} GB"
+    except:
+        ram_used = "Unavailable"
 
     # 🔹 PyTorch Generate
-    pt_output_text = ""
     pt_start = time.time()
     try:
         torch_inputs = tokenizer(prompt, return_tensors="pt")
@@ -46,9 +39,10 @@ def compare_outputs(prompt):
         pt_model.eval()
         with torch.no_grad():
             pt_outputs = pt_model.generate(**torch_inputs, max_new_tokens=50)
-        pt_output_text = tokenizer.decode(pt_outputs[0], skip_special_tokens=True)
+        pt_output_ids = pt_outputs[0].tolist()
+        pt_output_text = tokenizer.decode(pt_output_ids, skip_special_tokens=True)
+        pt_tokens = tokenizer.convert_ids_to_tokens(pt_output_ids)
         pt_time = time.time() - pt_start
-        summary_log.append(f"🧠 PyTorch output length: {pt_outputs.shape[1]} tokens | Time: {pt_time:.2f}s")
     finally:
         del pt_model
         gc.collect()
@@ -56,22 +50,49 @@ def compare_outputs(prompt):
             torch.cuda.empty_cache()
 
     # 🔹 ONNX Generate (Greedy)
-    ort_output_text = ""
     ort_start = time.time()
     ort_inputs = tokenizer(prompt, return_tensors="np")
     onnx_path = hf_hub_download(repo_id=HF_ONNX_REPO, filename=ONNX_MODEL_FILE)
     ort_session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
-    ort_output_text = greedy_decode_onnx(
-        ort_session, ort_inputs["input_ids"], ort_inputs["attention_mask"], max_new_tokens=50
-    )
+    ort_output_ids = []
+    generated = ort_inputs["input_ids"]
+    attention_mask = ort_inputs["attention_mask"]
+    for _ in range(50):
+        ort_outputs = ort_session.run(None, {
+            "input_ids": generated,
+            "attention_mask": attention_mask
+        })
+        next_token_logits = ort_outputs[0][:, -1, :]
+        next_token = np.argmax(next_token_logits, axis=-1).reshape(-1, 1)
+        ort_output_ids.append(next_token[0][0])
+        generated = np.concatenate((generated, next_token), axis=1)
+        attention_mask = np.concatenate((attention_mask, np.ones((1, 1), dtype=np.int64)), axis=1)
+        if next_token[0][0] == tokenizer.eos_token_id:
+            break
     ort_time = time.time() - ort_start
-    summary_log.append(f"⚙️ ONNX output length: {len(tokenizer(ort_output_text)['input_ids'])} tokens | Time: {ort_time:.2f}s")
+    ort_tokens = tokenizer.convert_ids_to_tokens(ort_inputs["input_ids"][0].tolist() + ort_output_ids)
+    ort_output_text = tokenizer.decode(ort_inputs["input_ids"][0].tolist() + ort_output_ids, skip_special_tokens=True)
 
-    # Final notes
-    summary_log.append(f"🧪 Tokenizer source: {tokenizer.name_or_path} | Vocab size: {tokenizer.vocab_size}")
-    summary_log.append("💡 Note: Future versions will include quantized ONNX (INT8) + Vitis AI support.")
+    # 📊 Summary
+    summary_log.append("| Model   | Tokens | Time (s) | Time/Token |")
+    summary_log.append("|---------|--------|----------|------------|")
+    summary_log.append(f"| PyTorch | {len(pt_tokens)} | {pt_time:.2f} | {pt_time / max(1, len(pt_tokens)):.4f} |")
+    summary_log.append(f"| ONNX    | {len(ort_tokens)} | {ort_time:.2f} | {ort_time / max(1, len(ort_tokens)):.4f} |")
+    summary_log.append(f"\n📦 RAM Used: {ram_used}")
+    summary_log.append(f"📚 Tokenizer: {tokenizer.name_or_path} | Vocab size: {tokenizer.vocab_size}")
+    summary_log.append("🛠️ Note: This ONNX export is FP32. INT8 + Vitis AI variants coming soon.")
 
-    return pt_output_text, ort_output_text, "\n".join(summary_log)
+    outputs = [pt_output_text, ort_output_text, "\n".join(summary_log)]
+
+    if show_tokens:
+        outputs += [
+            ", ".join(pt_tokens),
+            ", ".join(ort_tokens)
+        ]
+    else:
+        outputs += ["", ""]
+
+    return outputs
 
 example_prompts = [
     "Who was the first president of the United States?",
@@ -83,15 +104,20 @@ example_prompts = [
 
 iface = gr.Interface(
     fn=compare_outputs,
-    inputs=gr.Textbox(lines=2, placeholder="Enter a prompt..."),
+    inputs=[
+        gr.Textbox(lines=2, placeholder="Enter a prompt..."),
+        gr.Checkbox(label="Show Token IDs")
+    ],
     outputs=[
         gr.Textbox(label="PyTorch Output"),
         gr.Textbox(label="ONNX Output"),
-        gr.Textbox(label="Test Summary & Metadata")
+        gr.Textbox(label="Evaluation Summary"),
+        gr.Textbox(label="PyTorch Tokens"),
+        gr.Textbox(label="ONNX Tokens")
     ],
-    title="ONNX vs PyTorch (Full Output Comparison)",
-    description="Sequentially runs both models on your prompt and returns decoded output + metadata.",
-    examples=[[p] for p in example_prompts]
+    title="ONNX vs PyTorch (Full Output + Token Trace)",
+    description="Run both models on your prompt and compare output text, timing, and token traces. Sequential model loading avoids OOM.",
+    examples=[[p, False] for p in example_prompts]
 )
 
 iface.launch()