Update handler.py

handler.py

@@ -1,12 +1,10 @@
 # -*- coding: utf-8 -*-
 """
-PULSE ECG Handler - Demo-like sampling + no_stop (hard) + min_new_tokens auto
-- do_sample=True, temperature/top_p payload'dan
-- max_new_tokens: payload değeri (kırpma yok)
-- no_stop=True: stopping_criteria KAPALI + eos_token_id=None
-- no_stop=True ve min_new_tokens boşsa: otomatik min_new_tokens (uzun yanıt garantisi)
-- Tek görsel; IM_START/END otomatik; 3D/4D/5D tensör uyumlu
-- Post-format yok (demo davranışı)
+PULSE ECG Handler (demo-like streaming)
+- TextIteratorStreamer + skip_prompt=True (dilimleme yok; Step 1 korunur)
+- do_sample=True (demo davranışı), temperature/top_p payload'dan
+- Opsiyonel: no_stop, custom_stop, no_repeat_ngram_size, min_new_tokens
+- IM_START/END otomatik; 3D/4D/5D görüntü tensörü uyumlu; device/dtype eşleştirme
 """
 
 import os
@@ -15,20 +13,14 @@ import base64
 import hashlib
 import datetime
 from io import BytesIO
+from threading import Thread
+from typing import Optional, List
 
 import torch
 from PIL import Image
 import requests
 
-# --- Opsiyonel ---
-try:
-    import cv2
-    CV2_AVAILABLE = True
-except Exception:
-    CV2_AVAILABLE = False
-    print("Warning: OpenCV (cv2) not available; video is disabled.")
-
-# --- LLaVA ---
+# --- LLaVA / Transformers ---
 try:
     from llava.constants import (
         IMAGE_TOKEN_INDEX,
@@ -48,7 +40,14 @@ try:
     LLAVA_AVAILABLE = True
 except Exception as e:
     LLAVA_AVAILABLE = False
-    print(f"Warning: LLaVA modules not available: {e}")
+    print(f"[WARN] LLaVA modules not available: {e}")
+
+try:
+    from transformers import TextIteratorStreamer
+    TRANSFORMERS_AVAILABLE = True
+except Exception as e:
+    TRANSFORMERS_AVAILABLE = False
+    print(f"[WARN] transformers not available: {e}")
 
 # --- HF Hub (opsiyonel logging) ---
 try:
@@ -70,10 +69,9 @@ if HF_HUB_AVAILABLE and "HF_TOKEN" in os.environ:
         repo_name = ""
 
 LOGDIR = "./logs"
-VOTEDIR = "./votes"
 os.makedirs(LOGDIR, exist_ok=True)
-os.makedirs(VOTEDIR, exist_ok=True)
 
+# --- Global Model State ---
 tokenizer = None
 model = None
 image_processor = None
@@ -81,8 +79,10 @@ context_len = None
 args = None
 model_initialized = False
 
+# ----------------- Utilities -----------------
+
 def _safe_upload(path: str):
-    if api and repo_name and os.path.isfile(path):
+    if api and repo_name and path and os.path.isfile(path):
         try:
             api.upload_file(
                 path_or_fileobj=path,
@@ -100,6 +100,12 @@ def _conv_log_path():
     return p
 
 def load_image_any(image_input):
+    """
+    Desteklenen:
+      - URL (http/https)
+      - Yerel dosya yolu
+      - base64 (opsiyonel data URL prefix ile)
+    """
     if isinstance(image_input, str):
         s = image_input.strip()
         if s.startswith(("http://", "https://")):
@@ -108,6 +114,7 @@ def load_image_any(image_input):
             return Image.open(BytesIO(r.content)).convert("RGB")
         if os.path.exists(s):
             return Image.open(s).convert("RGB")
+        # base64
         if s.startswith("data:image"):
             s = s.split(",", 1)[1]
         raw = base64.b64decode(s)
@@ -134,72 +141,87 @@ def _wrap_image_token_if_needed(model_cfg) -> bool:
 def _build_prompt_and_ids(chatbot, user_text: str, device: torch.device):
     use_wrap = _wrap_image_token_if_needed(chatbot.model.config)
     if use_wrap:
+        # <im_start><image><im_end>\n + user text
         inp = f"{DEFAULT_IM_START_TOKEN}{DEFAULT_IMAGE_TOKEN}{DEFAULT_IM_END_TOKEN}\n{user_text}"
     else:
         inp = f"{DEFAULT_IMAGE_TOKEN}\n{user_text}"
+
     chatbot.conversation.append_message(chatbot.conversation.roles[0], inp)
     chatbot.conversation.append_message(chatbot.conversation.roles[1], None)
     prompt = chatbot.conversation.get_prompt()
+
     input_ids = tokenizer_image_token(
         prompt, chatbot.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
     ).unsqueeze(0).to(device)
     return prompt, input_ids
 
-def _stopping(chatbot, input_ids):
+def _stopping_keywords(chatbot, input_ids, extra: Optional[List[str]] = None):
     conv = chatbot.conversation
     stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
-    return KeywordsStoppingCriteria([stop_str], chatbot.tokenizer, input_ids)
+    keys = [stop_str]
+    if extra:
+        keys.extend([k for k in extra if isinstance(k, str) and k.strip()])
+    return KeywordsStoppingCriteria(keys, chatbot.tokenizer, input_ids)
+
+# ----------------- Core Generation -----------------
 
 def generate_response(
     message_text: str,
     image_input,
     *,
-    max_new_tokens: int = 4096,
-    temperature: float = 0.05,
-    top_p: float = 1.0,
-    repetition_penalty: float = 1.0,
-    conv_mode_override: str | None = None,
-    det_seed: int | None = None,
+    max_new_tokens: int = 1800,
+    min_new_tokens: Optional[int] = None,
+    temperature: float = 0.20,
+    top_p: float = 0.95,
+    repetition_penalty: float = 1.20,
+    no_repeat_ngram_size: Optional[int] = 6,
+    conv_mode_override: Optional[str] = None,
+    det_seed: Optional[int] = None,
     no_stop: bool = False,
-    min_new_tokens: int | None = None,   # otomatik atanabilir
+    custom_stop: Optional[List[str]] = None,
 ):
-    if not LLAVA_AVAILABLE:
-        return {"error": "LLaVA modules not available"}
+    if not (LLAVA_AVAILABLE and TRANSFORMERS_AVAILABLE):
+        return {"error": "Required libraries not available (llava/transformers)"}
     if not message_text or image_input is None:
         return {"error": "Both 'message' and 'image' are required"}
 
+    # Chat session (fresh conv each call, demo-like)
     chatbot = chat_manager.get_chatbot(args, args.model_path, tokenizer, model, image_processor, context_len)
     if conv_mode_override and conv_mode_override in conv_templates:
         chatbot.conversation = conv_templates[conv_mode_override].copy()
     else:
         chatbot.conversation = conv_templates[chatbot.conv_mode].copy()
 
+    # Load image
     try:
         pil_img = load_image_any(image_input)
     except Exception as e:
         return {"error": f"Failed to load image: {e}"}
 
-    # log
+    # Save image to logs (optional)
     img_hash, img_path = "NA", None
     try:
-        buf = BytesIO(); pil_img.save(buf, format="JPEG"); img_bytes = buf.getvalue()
-        img_hash = hashlib.md5(img_bytes).hexdigest()
+        buf = BytesIO(); pil_img.save(buf, format="JPEG"); raw = buf.getvalue()
+        img_hash = hashlib.md5(raw).hexdigest()
         t = datetime.datetime.now()
         img_path = os.path.join(LOGDIR, "serve_images", f"{t.year:04d}-{t.month:02d}-{t.day:02d}", f"{img_hash}.jpg")
         os.makedirs(os.path.dirname(img_path), exist_ok=True)
-        if not os.path.isfile(img_path): pil_img.save(img_path)
+        if not os.path.isfile(img_path):
+            pil_img.save(img_path)
     except Exception as e:
         print(f"[log] saving image failed: {e}")
 
-    # görüntü tensörü
+    # To device/dtype
     device = next(chatbot.model.parameters()).device
     dtype  = next(chatbot.model.parameters()).dtype
+
+    # Preprocess image -> tensor (support 3D/4D/5D)
     try:
         processed = process_images([pil_img], chatbot.image_processor, chatbot.model.config)
         if isinstance(processed, torch.Tensor):
             if processed.ndim == 3:   image_tensor = processed.unsqueeze(0)
             elif processed.ndim == 4: image_tensor = processed
-            elif processed.ndim == 5:
+            elif processed.ndim == 5: # (B,T,C,H,W) -> (B*T,C,H,W)
                 b,t,c,h,w = processed.shape
                 image_tensor = processed.reshape(b*t, c, h, w)
             else:
@@ -213,11 +235,16 @@ def generate_response(
     except Exception as e:
         return {"error": f"Image processing failed: {e}"}
 
-    # prompt & ids
+    # Prompt & ids
     _, input_ids = _build_prompt_and_ids(chatbot, message_text, device)
-    stopping = None if no_stop else _stopping(chatbot, input_ids)
 
-    # deterministik sample (opsiyonel)
+    # Stopping criteria
+    stopping = None if no_stop else _stopping_keywords(chatbot, input_ids, custom_stop)
+    eos_id = chatbot.tokenizer.eos_token_id
+    pad_id = chatbot.tokenizer.pad_token_id if chatbot.tokenizer.pad_token_id is not None else (eos_id if eos_id is not None else 0)
+    eos_for_gen = None if no_stop else eos_id
+
+    # Deterministic sampling (optional)
     if det_seed is not None:
         try:
             det_seed = int(det_seed)
@@ -228,18 +255,15 @@ def generate_response(
         except Exception:
             pass
 
-    # EOS/PAD
-    eos_id = tokenizer.eos_token_id
-    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else (eos_id if eos_id is not None else 0)
-    # no_stop=True ise eos'a göre durmayı tamamen kapat
-    if no_stop:
-        eos_for_gen = None
-    else:
-        eos_for_gen = eos_id
+    # Streamer (demo-like, avoids manual slicing)
+    streamer = TextIteratorStreamer(
+        chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
 
     gen_kwargs = dict(
         inputs=input_ids,
         images=image_tensor,
+        streamer=streamer,
         do_sample=True,
         temperature=float(temperature),
         top_p=float(top_p),
@@ -250,37 +274,38 @@ def generate_response(
         eos_token_id=eos_for_gen,
         length_penalty=1.0,
         early_stopping=False,
-        stopping_criteria=None if no_stop else [stopping],
+        stopping_criteria=None if no_stop else ([stopping] if stopping else None),
     )
 
-    # min_new_tokens otomatik (no_stop=True ve kullanıcı vermediyse)
-    if no_stop and (min_new_tokens is None):
+    if no_repeat_ngram_size:
         try:
-            req = int(max_new_tokens)
-            auto_min = max(300, min(req - 64, 1024))  # 300–1024 bandında güvenli
-            if auto_min > 0:
-                gen_kwargs["min_new_tokens"] = auto_min
+            n = int(no_repeat_ngram_size)
+            if n > 0:
+                gen_kwargs["no_repeat_ngram_size"] = n
         except Exception:
             pass
-    elif min_new_tokens is not None:
+
+    if min_new_tokens is not None:
         try:
             mn = int(min_new_tokens)
-            if mn > 0 and mn <= int(max_new_tokens):
+            if 1 <= mn <= int(max_new_tokens):
                 gen_kwargs["min_new_tokens"] = mn
         except Exception:
             pass
 
-    # generate
+    # Generate in a background thread; collect streamed tokens
     try:
-        with torch.no_grad():
-            outputs = chatbot.model.generate(**gen_kwargs)
-        gen = outputs[0][input_ids.shape[1]:]
-        text = tokenizer.decode(gen, skip_special_tokens=True)
+        t = Thread(target=chatbot.model.generate, kwargs=gen_kwargs)
+        t.start()
+        chunks = []
+        for piece in streamer:
+            chunks.append(piece)
+        text = "".join(chunks)
         chatbot.conversation.messages[-1][-1] = text
     except Exception as e:
         return {"error": f"Generation failed: {e}"}
 
-    # log
+    # Log
     try:
         row = {
             "time": datetime.datetime.now().isoformat(),
@@ -292,16 +317,16 @@ def generate_response(
         }
         with open(_conv_log_path(), "a", encoding="utf-8") as f:
             f.write(json.dumps(row, ensure_ascii=False) + "\n")
-        _safe_upload(_conv_log_path());  _safe_upload(img_path or "")
+        _safe_upload(_conv_log_path()); _safe_upload(img_path or "")
     except Exception as e:
         print(f"[log] failed: {e}")
 
     return {"status": "success", "response": text, "conversation_id": id(chatbot.conversation)}
 
-# -------- API --------
+# ----------------- Public API -----------------
 
 def query(payload: dict):
-    """HF Endpoint entry (demo-like)"""
+    """HF Endpoint entry (demo-like streaming)"""
     global model_initialized, tokenizer, model, image_processor, context_len, args
     if not model_initialized:
         if not initialize_model():
@@ -314,32 +339,43 @@ def query(payload: dict):
         if not message.strip(): return {"error": "Missing 'message' text"}
         if image is None:       return {"error": "Missing 'image'. Use 'image', 'image_url', or 'img'."}
 
-        max_new_tokens      = int(payload.get("max_output_tokens", payload.get("max_new_tokens", payload.get("max_tokens", 4096))))
-        temperature         = float(payload.get("temperature", 0.05))
-        top_p               = float(payload.get("top_p", 1.0))
-        repetition_penalty  = float(payload.get("repetition_penalty", 1.0))
+        # Demo-like knobs
+        max_new_tokens      = int(payload.get("max_output_tokens", payload.get("max_new_tokens", payload.get("max_tokens", 1800))))
+        min_new_tokens      = payload.get("min_new_tokens", None)
+        if min_new_tokens is not None:
+            try: min_new_tokens = int(min_new_tokens)
+            except Exception: min_new_tokens = None
+
+        temperature         = float(payload.get("temperature", 0.20))
+        top_p               = float(payload.get("top_p", 0.95))
+        repetition_penalty  = float(payload.get("repetition_penalty", 1.20))
+        no_repeat_ngram     = payload.get("no_repeat_ngram_size", 6)
+        try:
+            no_repeat_ngram = int(no_repeat_ngram) if no_repeat_ngram is not None else None
+        except Exception:
+            no_repeat_ngram = None
+
         conv_mode_override  = payload.get("conv_mode", None)
         det_seed            = payload.get("det_seed", None)
         if det_seed is not None:
             try: det_seed = int(det_seed)
             except Exception: det_seed = None
         no_stop             = bool(payload.get("no_stop", False))
-        mnt                 = payload.get("min_new_tokens", None)
-        if mnt is not None:
-            try: mnt = int(mnt)
-            except Exception: mnt = None
+        custom_stop         = payload.get("custom_stop", None)
 
         return generate_response(
             message_text=message,
             image_input=image,
             max_new_tokens=max_new_tokens,
+            min_new_tokens=min_new_tokens,
             temperature=temperature,
             top_p=top_p,
             repetition_penalty=repetition_penalty,
+            no_repeat_ngram_size=no_repeat_ngram,
             conv_mode_override=conv_mode_override,
             det_seed=det_seed,
             no_stop=no_stop,
-            min_new_tokens=mnt,
+            custom_stop=custom_stop,
         )
     except Exception as e:
         return {"error": f"Query failed: {e}"}
@@ -350,7 +386,7 @@ def health_check():
         "model_initialized": model_initialized,
         "cuda_available": torch.cuda.is_available(),
         "llava_available": LLAVA_AVAILABLE,
-        "cv2_available": CV2_AVAILABLE,
+        "transformers_available": TRANSFORMERS_AVAILABLE,
     }
 
 def get_model_info():
@@ -362,7 +398,7 @@ def get_model_info():
         "device": str(next(model.parameters()).device) if model else "Unknown",
     }
 
-# -------- init --------
+# ----------------- Init & Session -----------------
 
 class _Args:
     def __init__(self):
@@ -370,19 +406,19 @@ class _Args:
         self.model_base = None
         self.num_gpus   = int(os.getenv("NUM_GPUS", "1"))
         self.conv_mode  = None
-        self.max_new_tokens = int(os.getenv("MAX_NEW_TOKENS", "4096"))
+        self.max_new_tokens = int(os.getenv("MAX_NEW_TOKENS", "1800"))
         self.num_frames = 16
         self.load_8bit  = bool(int(os.getenv("LOAD_8BIT", "0")))
         self.load_4bit  = bool(int(os.getenv("LOAD_4BIT", "0")))
         self.debug      = bool(int(os.getenv("DEBUG", "0")))
 
 class InferenceDemo:
-    def __init__(self, args, model_path, tokenizer, model, image_processor, context_len):
+    def __init__(self, args, model_path, tokenizer_, model_, image_processor_, context_len_):
         if not LLAVA_AVAILABLE:
             raise ImportError("LLaVA modules not available")
         disable_torch_init()
         self.tokenizer, self.model, self.image_processor, self.context_len = (
-            tokenizer, model, image_processor, context_len
+            tokenizer_, model_, image_processor_, context_len_
         )
         conv_mode_auto = _guess_conv_mode(model_path)
         self.conv_mode = args.conv_mode if args.conv_mode else conv_mode_auto
@@ -409,7 +445,7 @@ chat_manager = ChatSessionManager()
 def initialize_model():
     global tokenizer, model, image_processor, context_len, args
     if not LLAVA_AVAILABLE:
-        print("LLaVA not available; cannot init.")
+        print("[init] LLaVA not available; cannot init.")
         return False
     try:
         args = _Args()
@@ -430,7 +466,10 @@ def initialize_model():
         print(f"[init] failed: {e}")
         return False
 
+# ----------------- HF EndpointHandler -----------------
+
 class EndpointHandler:
+    """Hugging Face Endpoint uyumlu sınıf"""
     def __init__(self, model_dir):
         self.model_dir = model_dir
         print(f"EndpointHandler initialized with model_dir: {model_dir}")