Update app.py

app.py

@@ -1,13 +1,11 @@
 import os
 import gradio as gr
 import torch
-import numpy as np
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 from concurrent.futures import ThreadPoolExecutor
 import logging
 import time
-import queue
 import threading
 
 logging.basicConfig(level=logging.INFO)
@@ -16,21 +14,18 @@ logger = logging.getLogger(__name__)
 MODEL_ID = "YoussefElsafi/PlayerAI-1.2B-v1.5"
 
 # ── CPU Threading setup FIRST (before any torch ops) ─────────────────────────
-# On CPU, torch uses OpenMP/MKL threads internally.
-# Too many = thrashing. Formula: physical_cores / expected_concurrent_users
 PHYSICAL_CORES  = os.cpu_count() or 2
-MAX_CONCURRENT  = 2          # how many generations run truly in parallel
+MAX_CONCURRENT  = 2
 THREADS_PER_GEN = max(1, PHYSICAL_CORES // MAX_CONCURRENT)
 
 torch.set_num_threads(THREADS_PER_GEN)
 torch.set_num_interop_threads(2)
 
-# Force MKL/OpenBLAS to respect our thread budget too
-os.environ["OMP_NUM_THREADS"]       = str(THREADS_PER_GEN)
-os.environ["MKL_NUM_THREADS"]       = str(THREADS_PER_GEN)
-os.environ["OPENBLAS_NUM_THREADS"]  = str(THREADS_PER_GEN)
-os.environ["VECLIB_MAXIMUM_THREADS"]= str(THREADS_PER_GEN)
-os.environ["NUMEXPR_NUM_THREADS"]   = str(THREADS_PER_GEN)
+os.environ["OMP_NUM_THREADS"]        = str(THREADS_PER_GEN)
+os.environ["MKL_NUM_THREADS"]        = str(THREADS_PER_GEN)
+os.environ["OPENBLAS_NUM_THREADS"]   = str(THREADS_PER_GEN)
+os.environ["VECLIB_MAXIMUM_THREADS"] = str(THREADS_PER_GEN)
+os.environ["NUMEXPR_NUM_THREADS"]    = str(THREADS_PER_GEN)
 
 logger.info(f"CPU cores: {PHYSICAL_CORES} | "
             f"Concurrent gens: {MAX_CONCURRENT} | "
@@ -40,45 +35,52 @@ logger.info(f"CPU cores: {PHYSICAL_CORES} | "
 print("Loading model...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-# On CPU: bfloat16 is FASTER than float32 on modern x86 (AVX-512 bf16 support)
-# Falls back gracefully if not supported
-if torch.cuda.is_available():
-    DTYPE  = torch.float16
-    DEVICE = "cuda"
-elif hasattr(torch, 'bfloat16') and torch.backends.cpu.is_avx512_bf16_supported():
-    DTYPE  = torch.bfloat16
-    DEVICE = "cpu"
-    logger.info("AVX-512 BF16 detected — using bfloat16")
-else:
-    DTYPE  = torch.float32
-    DEVICE = "cpu"
+# ── Safe dtype detection (works on ALL torch versions) ────────────────────────
+def _pick_dtype() -> torch.dtype:
+    if torch.cuda.is_available():
+        logger.info("Device: CUDA — using float16")
+        return torch.float16
+
+    # bfloat16 on CPU: supported natively on Ampere+ and some Intel CPUs.
+    # We probe by actually trying a small bf16 op instead of relying on
+    # version-specific API calls that may not exist.
+    try:
+        _ = torch.zeros(4, dtype=torch.bfloat16) @ torch.zeros(4, dtype=torch.bfloat16)
+        logger.info("Device: CPU — bfloat16 probe passed, using bfloat16")
+        return torch.bfloat16
+    except Exception:
+        logger.info("Device: CPU — bfloat16 not supported, using float32")
+        return torch.float32
+
+DTYPE  = _pick_dtype()
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype        = DTYPE,
-    device_map         = "cpu" if DEVICE == "cpu" else "auto",
-    trust_remote_code  = True,
-    low_cpu_mem_usage  = True,
+    torch_dtype       = DTYPE,
+    device_map        = "auto" if DEVICE == "cuda" else "cpu",
+    trust_remote_code = True,
+    low_cpu_mem_usage = True,
 )
 model.eval()
 
-# ── INT8 dynamic quantization (biggest CPU speedup: ~2x faster, ~2x less RAM) ─
-# Works on Linear layers — the bulk of transformer compute
-try:
-    model = torch.quantization.quantize_dynamic(
-        model,
-        {torch.nn.Linear},
-        dtype=torch.qint8,
-    )
-    logger.info("INT8 dynamic quantization applied ✓")
-except Exception as e:
-    logger.warning(f"Quantization skipped: {e}")
+# ── INT8 dynamic quantization ─────────────────────────────────────────────────
+# ~2× faster inference, ~50% less RAM. Only apply on CPU.
+if DEVICE == "cpu":
+    try:
+        model = torch.quantization.quantize_dynamic(
+            model,
+            {torch.nn.Linear},
+            dtype=torch.qint8,
+        )
+        logger.info("INT8 dynamic quantization applied ✓")
+    except Exception as e:
+        logger.warning(f"Quantization skipped: {e}")
 
-# ── torch.compile (PyTorch 2.0+) ──────────────────────────────────────────────
-# "default" mode works better than "reduce-overhead" on CPU
+# ── torch.compile ─────────────────────────────────────────────────────────────
 if hasattr(torch, "compile"):
     try:
-        model = torch.compile(model, mode="default", backend="inductor")
+        model = torch.compile(model, mode="default")
         logger.info("torch.compile applied ✓")
     except Exception as e:
         logger.warning(f"torch.compile skipped: {e}")
@@ -86,30 +88,19 @@ if hasattr(torch, "compile"):
 print("Model loaded!")
 
 # ── Generation config ─────────────────────────────────────────────────────────
-# Shorter = faster. For chat, 80 tokens is plenty.
 GEN_CONFIG = dict(
-    max_new_tokens  = 80,       # was 120 — shaves ~33% off latency
-    do_sample       = True,
-    temperature     = 0.85,
-    top_p           = 0.9,
-    # repetition_penalty speeds up sampling by pruning logits
+    max_new_tokens     = 80,
+    do_sample          = True,
+    temperature        = 0.85,
+    top_p              = 0.9,
     repetition_penalty = 1.1,
-    pad_token_id    = tokenizer.eos_token_id,
-    # Use cache — critical for autoregressive speed
-    use_cache       = True,
+    pad_token_id       = tokenizer.eos_token_id,
+    use_cache          = True,
 )
 
-# ── Thread pool ───────────────────────────────────────────────────────────────
-executor = ThreadPoolExecutor(
-    max_workers      = MAX_CONCURRENT,
-    thread_name_prefix = "gen",
-)
-
-# ── Request queue + priority ───────────────────────────────────────────────────
-# Caps in-flight requests so we don't OOM with 16GB RAM
-# 1.2B model in int8 ≈ 1.2GB, each active KV cache ≈ 200MB
-# So 2 concurrent is safe, queue up to 8 waiting
-MAX_QUEUE    = 8
+# ── Thread pool + semaphores ──────────────────────────────────────────────────
+executor      = ThreadPoolExecutor(max_workers=MAX_CONCURRENT, thread_name_prefix="gen")
+MAX_QUEUE     = 8
 gen_semaphore = threading.Semaphore(MAX_CONCURRENT)
 queue_counter = threading.Semaphore(MAX_QUEUE)
 
@@ -128,27 +119,14 @@ SYSTEM_PROMPT = (
     "Use lowercase, slang, and act like a real person chatting."
 )
 
-# ── Prompt cache ───────────────────────────────────────────────────────────────
-# System prompt tokenized once and reused — saves ~5ms per request
-_SYSTEM_MESSAGES = [{"role": "system", "content": SYSTEM_PROMPT}]
-try:
-    _SYSTEM_PROMPT_PREFIX = tokenizer.apply_chat_template(
-        _SYSTEM_MESSAGES,
-        tokenize=False,
-        add_generation_prompt=False,
-    )
-except Exception:
-    _SYSTEM_PROMPT_PREFIX = ""
-
 # ── Core generation ────────────────────────────────────────────────────────────
 def _run_generation(prompt: str, streamer: TextIteratorStreamer) -> None:
     inputs = tokenizer(
         prompt,
-        return_tensors    = "pt",
-        truncation        = True,
-        max_length        = 512,   # hard cap — prevents runaway long contexts
+        return_tensors = "pt",
+        truncation     = True,
+        max_length     = 512,
     )
-    # No .to(device) needed — model is already on cpu
     with torch.inference_mode():
         model.generate(
             input_ids      = inputs.input_ids,
@@ -157,23 +135,19 @@ def _run_generation(prompt: str, streamer: TextIteratorStreamer) -> None:
             **GEN_CONFIG,
         )
 
-
-# ── Streaming response ────────────────────────────────────────────────────────
+# ── Streaming response ─────────────────────────────────────────────────────────
 def respond(message: str, history: list, api_key: str):
     if not validate_key(api_key):
         yield "[error] unauthorized"
         return
 
-    # Queue guard — reject if backlog is full
     if not queue_counter.acquire(blocking=False):
         yield "[error] server busy — try again in a moment"
         return
 
     try:
-        # Build prompt — keep history short to limit context length
-        # Long contexts = slow prefill on CPU
-        MAX_HISTORY = 4   # last 4 turns only
-        trimmed_history = history[-MAX_HISTORY:] if len(history) > MAX_HISTORY else history
+        MAX_HISTORY     = 4
+        trimmed_history = history[-MAX_HISTORY:]
 
         messages = [{"role": "system", "content": SYSTEM_PROMPT}]
         messages.extend(trimmed_history)
@@ -182,7 +156,7 @@ def respond(message: str, history: list, api_key: str):
         try:
             prompt = tokenizer.apply_chat_template(
                 messages,
-                tokenize           = False,
+                tokenize              = False,
                 add_generation_prompt = True,
             )
         except Exception as e:
@@ -191,12 +165,11 @@ def respond(message: str, history: list, api_key: str):
 
         streamer = TextIteratorStreamer(
             tokenizer,
-            skip_prompt        = True,
-            skip_special_tokens= True,
-            timeout            = 60.0,
+            skip_prompt         = True,
+            skip_special_tokens = True,
+            timeout             = 60.0,
         )
 
-        # Acquire generation slot (blocks if MAX_CONCURRENT already running)
         gen_semaphore.acquire()
         try:
             future = executor.submit(_run_generation, prompt, streamer)
@@ -207,7 +180,6 @@ def respond(message: str, history: list, api_key: str):
                 yield response
 
             future.result(timeout=10)
-
         finally:
             gen_semaphore.release()
 
@@ -215,23 +187,25 @@ def respond(message: str, history: list, api_key: str):
         queue_counter.release()
 
 
-# ── Warm-up ────────────────────────────────────────────────────────────────────
-# Run one dummy generation at startup so the first real user isn't slow.
-# torch.compile needs a first pass to JIT-compile the graph.
+# ── Warmup ────────────────────────────────────────────────────────────────────
 def _warmup():
     logger.info("Running warmup generation...")
     t0 = time.time()
     try:
         dummy_streamer = TextIteratorStreamer(
-            tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=120.0
+            tokenizer,
+            skip_prompt         = True,
+            skip_special_tokens = True,
+            timeout             = 120.0,
         )
         dummy_prompt = tokenizer.apply_chat_template(
             [{"role": "system", "content": SYSTEM_PROMPT},
              {"role": "user",   "content": "hi"}],
-            tokenize=False,
-            add_generation_prompt=True,
+            tokenize              = False,
+            add_generation_prompt = True,
         )
-        inputs = tokenizer(dummy_prompt, return_tensors="pt", truncation=True, max_length=256)
+        inputs     = tokenizer(dummy_prompt, return_tensors="pt",
+                               truncation=True, max_length=256)
         warmup_cfg = {**GEN_CONFIG, "max_new_tokens": 20}
 
         def _gen():
@@ -248,11 +222,10 @@ def _warmup():
         for _ in dummy_streamer:
             pass
         t.join()
-        logger.info(f"Warmup done in {time.time()-t0:.2f}s ✓")
+        logger.info(f"Warmup done in {time.time() - t0:.2f}s ✓")
     except Exception as e:
         logger.warning(f"Warmup failed (non-fatal): {e}")
 
-# Run warmup in background so server starts immediately
 Thread(target=_warmup, daemon=True).start()
 
 
@@ -312,12 +285,12 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:
         stored_key = gr.Textbox(value="", visible=False)
 
         chatbot = gr.Chatbot(
-            type             = "messages",
-            avatar_images    = (USER_AVATAR, BOT_AVATAR),
-            height           = 500,
-            show_label       = False,
-            bubble_full_width= False,
-            show_copy_button = False,
+            type              = "messages",
+            avatar_images     = (USER_AVATAR, BOT_AVATAR),
+            height            = 500,
+            show_label        = False,
+            bubble_full_width = False,
+            show_copy_button  = False,
         )
 
         gr.ChatInterface(