feat: add smoke_test and generation metrics (latency/token counts)

app.py

@@ -11,6 +11,7 @@ from transformers.pipelines import pipeline
 from transformers import AutoTokenizer
 import torch
 import importlib
+import time
 
 
 # ---------------- Configuration ----------------
@@ -44,6 +45,15 @@ LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
 logging.basicConfig(level=LOG_LEVEL, format="%(asctime)s %(levelname)s %(message)s")
 logger = logging.getLogger("zamai-app")
 
+# Metrics storage for last real generation
+LAST_METRICS = {
+    "latency_sec": None,
+    "input_tokens": None,
+    "output_tokens": None,
+    "num_sequences": None,
+    "mode": None,
+}
+
 
 # ---------------- Utilities ----------------
 SAMPLE_INSTRUCTIONS = [
@@ -165,6 +175,7 @@ def predict(instruction: str,
 
     allowed_keys = {"max_new_tokens", "num_beams", "do_sample", "temperature", "top_p", "num_return_sequences"}
 
+    start = time.time()
     try:
         gen = get_generator()
         raw_kwargs = {
@@ -189,9 +200,28 @@ def predict(instruction: str,
             if text:
                 texts.append(text)
         if not texts:
+            LAST_METRICS.update({
+                "latency_sec": round(time.time() - start, 3),
+                "input_tokens": None,
+                "output_tokens": 0,
+                "num_sequences": 0,
+                "mode": active_mode,
+            })
             return f"### Prompt\n\n````\n{prompt}\n````\n\n### Output\n\n⚠️ No response generated."
         joined = "\n\n---\n\n".join(texts)
-        return f"### Prompt\n\n````\n{prompt}\n````\n\n### Output\n\n{joined}"
+
+        # Basic token counting via whitespace split (approximate)
+        input_tokens = len(prompt.split())
+        output_tokens = sum(len(t.split()) for t in texts)
+        LAST_METRICS.update({
+            "latency_sec": round(time.time() - start, 3),
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "num_sequences": len(texts),
+            "mode": active_mode,
+        })
+        metrics_md = f"\n\n### Metrics\n- Latency: {LAST_METRICS['latency_sec']}s\n- Input tokens (approx): {input_tokens}\n- Output tokens (approx): {output_tokens}\n- Sequences: {len(texts)}"
+        return f"### Prompt\n\n````\n{prompt}\n````\n\n### Output\n\n{joined}{metrics_md}"
     except Exception as e:
         logger.exception("Generation failed: %s", e)
         return f"⚠️ Generation failed: {e}"
@@ -245,7 +275,11 @@ def build_ui():
         instruction_dropdown.change(lambda x: x, inputs=instruction_dropdown, outputs=instruction_textbox)
 
         def refresh():
-            return f"**Device:** {'GPU' if _detect_device() != -1 else 'CPU'} | **Offline:** {os.getenv('HF_HUB_OFFLINE','0')} | **Env Mode:** {ECHO_MODE}"
+            base = f"**Device:** {'GPU' if _detect_device() != -1 else 'CPU'} | **Offline:** {os.getenv('HF_HUB_OFFLINE','0')} | **Env Mode:** {ECHO_MODE}"
+            if LAST_METRICS.get('latency_sec') is not None:
+                base += (f"<br>**Last Gen:** latency={LAST_METRICS['latency_sec']}s, "
+                         f"in≈{LAST_METRICS['input_tokens']}, out≈{LAST_METRICS['output_tokens']}, seqs={LAST_METRICS['num_sequences']}")
+            return base
 
         refresh_status.click(fn=refresh, inputs=None, outputs=status_box)
 

smoke_test.py

@@ -0,0 +1,21 @@
+import time
+from app import predict
+
+# Basic smoke tests for each mode.
+# Note: Real mode will load the model weights; keep max_new_tokens small.
+
+def run():
+    instruction = "ازموینه"  # Pashto for test
+    print("=== Echo Mode ===")
+    print(predict(instruction, "", 8, 2, True, 1.0, 0.9, 1, "echo"))
+    print("\n=== Useless Mode ===")
+    print(predict(instruction, "", 8, 2, True, 1.0, 0.9, 1, "useless"))
+    print("\n=== Real Mode (off) ===")
+    t0 = time.time()
+    out = predict(instruction, "", 8, 2, True, 1.0, 0.9, 1, "off")
+    dt = time.time() - t0
+    print(out)
+    print(f"\n[Latency real mode: {dt:.2f}s]")
+
+if __name__ == "__main__":
+    run()