Update handler.py

handler.py

@@ -1,10 +1,12 @@
 # -*- coding: utf-8 -*-
 """
-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
+PULSE ECG Handler (demo-like streaming, stable & clean)
+- TextIteratorStreamer + skip_prompt=True  → baş kesilmesi yok (Step 1 korunur)
+- do_sample=True (demo davranışı), temperature/top_p payload’dan
+- Anti-tekrar: no_repeat_ngram_size + repetition_penalty
+- Opsiyonel: custom_stop (örn. "END OF REPORT") → çıktı sonunda trim
+- Deterministik mod: aynı görüntü+mesaj için aynı seed (deterministic=True)
+- Görsel tensörü 3D/4D/5D uyumlu; device/dtype eşleştirme
 """
 
 import os
@@ -14,13 +16,13 @@ import hashlib
 import datetime
 from io import BytesIO
 from threading import Thread
-from typing import Optional, List
+from typing import Optional, List, Union
 
 import torch
 from PIL import Image
 import requests
 
-# --- LLaVA / Transformers ---
+# ---------- LLaVA & Transformers ----------
 try:
     from llava.constants import (
         IMAGE_TOKEN_INDEX,
@@ -34,7 +36,6 @@ try:
         tokenizer_image_token,
         process_images,
         get_model_name_from_path,
-        KeywordsStoppingCriteria,
     )
     from llava.utils import disable_torch_init
     LLAVA_AVAILABLE = True
@@ -49,7 +50,7 @@ except Exception as e:
     TRANSFORMERS_AVAILABLE = False
     print(f"[WARN] transformers not available: {e}")
 
-# --- HF Hub (opsiyonel logging) ---
+# ---------- HF Hub (opsiyonel logging) ----------
 try:
     from huggingface_hub import HfApi, login
     HF_HUB_AVAILABLE = True
@@ -71,7 +72,7 @@ if HF_HUB_AVAILABLE and "HF_TOKEN" in os.environ:
 LOGDIR = "./logs"
 os.makedirs(LOGDIR, exist_ok=True)
 
-# --- Global Model State ---
+# ---------- Global Model State ----------
 tokenizer = None
 model = None
 image_processor = None
@@ -79,7 +80,8 @@ context_len = None
 args = None
 model_initialized = False
 
-# ----------------- Utilities -----------------
+
+# ======================== Utilities ========================
 
 def _safe_upload(path: str):
     if api and repo_name and path and os.path.isfile(path):
@@ -93,18 +95,19 @@ def _safe_upload(path: str):
         except Exception as e:
             print(f"[upload] failed for {path}: {e}")
 
-def _conv_log_path():
+def _conv_log_path() -> str:
     t = datetime.datetime.now()
     p = os.path.join(LOGDIR, f"{t.year:04d}-{t.month:02d}-{t.day:02d}-user_conv.json")
     os.makedirs(os.path.dirname(p), exist_ok=True)
     return p
 
-def load_image_any(image_input):
+def load_image_any(image_input: Union[str, dict]) -> Image.Image:
     """
     Desteklenen:
       - URL (http/https)
       - Yerel dosya yolu
       - base64 (opsiyonel data URL prefix ile)
+      - {"image": <base64|dataurl>}
     """
     if isinstance(image_input, str):
         s = image_input.strip()
@@ -114,15 +117,16 @@ 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
+        # base64 (dataurl olabilir)
         if s.startswith("data:image"):
             s = s.split(",", 1)[1]
         raw = base64.b64decode(s)
         return Image.open(BytesIO(raw)).convert("RGB")
-    elif isinstance(image_input, dict) and "image" in image_input:
+
+    if isinstance(image_input, dict) and "image" in image_input:
         return load_image_any(image_input["image"])
-    else:
-        raise ValueError("Unsupported image input format")
+
+    raise ValueError("Unsupported image input format")
 
 def _guess_conv_mode(model_path: str) -> str:
     name = get_model_name_from_path(model_path).lower()
@@ -139,6 +143,7 @@ def _wrap_image_token_if_needed(model_cfg) -> bool:
         return False
 
 def _build_prompt_and_ids(chatbot, user_text: str, device: torch.device):
+    # Demo gibi: <image> token + text (IM_START/END gerekiyorsa sar)
     use_wrap = _wrap_image_token_if_needed(chatbot.model.config)
     if use_wrap:
         # <im_start><image><im_end>\n + user text
@@ -155,50 +160,50 @@ def _build_prompt_and_ids(chatbot, user_text: str, device: torch.device):
     ).unsqueeze(0).to(device)
     return prompt, 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
-    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)
+def _stable_seed_from(image_hash: str, message_text: str) -> int:
+    """Aynı resim+mesaj için aynı seed (deterministik örnekleme)"""
+    h = hashlib.md5((image_hash + "||" + message_text).encode("utf-8")).digest()
+    # 32-bit pozitif int
+    return int.from_bytes(h[:4], "big", signed=False)
 
-# ----------------- Core Generation -----------------
+
+# ======================== Core Generation ========================
 
 def generate_response(
     message_text: str,
     image_input,
     *,
     max_new_tokens: int = 1800,
-    min_new_tokens: Optional[int] = None,
+    min_new_tokens: Optional[int] = 700,
     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,
-    custom_stop: Optional[List[str]] = None,
+    deterministic: bool = False,  # True → do_sample=False (tam deterministik)
+    det_seed: Optional[int] = None,  # verilirse sabit seed
+    custom_stop: Optional[List[str]] = None,  # ["END OF REPORT"] gibi
+    no_stop: bool = False,  # True → eos/stop yok (önerilmez)
 ):
     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)
+    # Chat oturumu (her çağrıda taze template; demo benzeri)
     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
+    # Görseli yükle
     try:
         pil_img = load_image_any(image_input)
     except Exception as e:
         return {"error": f"Failed to load image: {e}"}
 
-    # Save image to logs (optional)
+    # Log için kaydet (hash + path)
     img_hash, img_path = "NA", None
     try:
         buf = BytesIO(); pil_img.save(buf, format="JPEG"); raw = buf.getvalue()
@@ -211,17 +216,17 @@ def generate_response(
     except Exception as e:
         print(f"[log] saving image failed: {e}")
 
-    # To device/dtype
+    # Cihaza/dtype’a taşı
     device = next(chatbot.model.parameters()).device
     dtype  = next(chatbot.model.parameters()).dtype
 
-    # Preprocess image -> tensor (support 3D/4D/5D)
+    # Görüntü ön-işleme → tensör (3D/4D/5D destek)
     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: # (B,T,C,H,W) -> (B*T,C,H,W)
+            if processed.ndim == 3:   image_tensor = processed.unsqueeze(0)     # (1,C,H,W)
+            elif processed.ndim == 4: image_tensor = processed                  # (B,C,H,W)
+            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:
@@ -238,33 +243,42 @@ def generate_response(
     # Prompt & ids
     _, input_ids = _build_prompt_and_ids(chatbot, message_text, device)
 
-    # 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)
+    # Seed ayarı
     if det_seed is not None:
         try:
-            det_seed = int(det_seed)
-            torch.manual_seed(det_seed)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed(det_seed)
-                torch.cuda.manual_seed_all(det_seed)
+            s = int(det_seed)
         except Exception:
-            pass
+            s = None
+    elif deterministic:
+        s = _stable_seed_from(img_hash, message_text)
+    else:
+        # Deterministik örnekleme istiyorsan; aynı girdide aynı sonuç için stabil seed de kullanabiliriz
+        s = _stable_seed_from(img_hash, message_text)
+
+    if s is not None:
+        torch.manual_seed(s)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(s)
+            torch.cuda.manual_seed_all(s)
+
+    # Stopping / EOS
+    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
 
-    # Streamer (demo-like, avoids manual slicing)
+    # Streamer (demo gibi; manuel dilimleme yok → Step 1 korunur)
     streamer = TextIteratorStreamer(
         chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True
     )
 
+    # do_sample: demo gibi (True). deterministic=True ise greedy’ye geç
+    do_sample = not deterministic
+
     gen_kwargs = dict(
         inputs=input_ids,
         images=image_tensor,
         streamer=streamer,
-        do_sample=True,
+        do_sample=do_sample,
         temperature=float(temperature),
         top_p=float(top_p),
         repetition_penalty=float(repetition_penalty),
@@ -274,7 +288,7 @@ def generate_response(
         eos_token_id=eos_for_gen,
         length_penalty=1.0,
         early_stopping=False,
-        stopping_criteria=None if no_stop else ([stopping] if stopping else None),
+        # stopping_criteria vermiyoruz → LLaVA'daki KeywordsStoppingCriteria hatalarından kaçınmak için
     )
 
     if no_repeat_ngram_size:
@@ -293,14 +307,24 @@ def generate_response(
         except Exception:
             pass
 
-    # Generate in a background thread; collect streamed tokens
+    # Üretim (arka thread) + stream toplama
     try:
         t = Thread(target=chatbot.model.generate, kwargs=gen_kwargs)
         t.start()
-        chunks = []
+        chunks: List[str] = []
         for piece in streamer:
             chunks.append(piece)
         text = "".join(chunks)
+        # custom_stop varsa çıktıdan itibaren kırp
+        if custom_stop:
+            if isinstance(custom_stop, str):
+                custom_stop = [custom_stop]
+            for tag in custom_stop:
+                if isinstance(tag, str) and tag:
+                    idx = text.find(tag)
+                    if idx != -1:
+                        text = text[:idx].rstrip()
+                        break
         chatbot.conversation.messages[-1][-1] = text
     except Exception as e:
         return {"error": f"Generation failed: {e}"}
@@ -323,7 +347,8 @@ def generate_response(
 
     return {"status": "success", "response": text, "conversation_id": id(chatbot.conversation)}
 
-# ----------------- Public API -----------------
+
+# ======================== Public API ========================
 
 def query(payload: dict):
     """HF Endpoint entry (demo-like streaming)"""
@@ -339,12 +364,13 @@ 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'."}
 
-        # Demo-like knobs
+        # Demo-like varsayılanlar
         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
+        min_new_tokens      = payload.get("min_new_tokens", 700)
+        try:
+            min_new_tokens = int(min_new_tokens) if min_new_tokens is not None else None
+        except Exception:
+            min_new_tokens = None
 
         temperature         = float(payload.get("temperature", 0.20))
         top_p               = float(payload.get("top_p", 0.95))
@@ -356,12 +382,14 @@ def query(payload: dict):
             no_repeat_ngram = None
 
         conv_mode_override  = payload.get("conv_mode", None)
+        deterministic       = bool(payload.get("deterministic", False))
         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))
+
         custom_stop         = payload.get("custom_stop", None)
+        no_stop             = bool(payload.get("no_stop", False))  # genelde False kalsın
 
         return generate_response(
             message_text=message,
@@ -373,9 +401,10 @@ def query(payload: dict):
             repetition_penalty=repetition_penalty,
             no_repeat_ngram_size=no_repeat_ngram,
             conv_mode_override=conv_mode_override,
+            deterministic=deterministic,
             det_seed=det_seed,
-            no_stop=no_stop,
             custom_stop=custom_stop,
+            no_stop=no_stop,
         )
     except Exception as e:
         return {"error": f"Query failed: {e}"}
@@ -398,7 +427,8 @@ def get_model_info():
         "device": str(next(model.parameters()).device) if model else "Unknown",
     }
 
-# ----------------- Init & Session -----------------
+
+# ======================== Init & Session ========================
 
 class _Args:
     def __init__(self):
@@ -409,6 +439,7 @@ class _Args:
         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")))
+        # 4bit/8bit hız için açık bırakılabilir; accelerate devicemap kullanıyorsanız .to(cuda) gerekmez
         self.load_4bit  = bool(int(os.getenv("LOAD_4BIT", "0")))
         self.debug      = bool(int(os.getenv("DEBUG", "0")))
 
@@ -450,23 +481,32 @@ def initialize_model():
     try:
         args = _Args()
         model_name = get_model_name_from_path(args.model_path)
-        tokenizer, model, image_processor, context_len = load_pretrained_model(
+        tokenizer_, model_, image_processor_, context_len_ = load_pretrained_model(
             args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit
         )
+        # Device
         try:
-            _ = next(model.parameters()).device
+            _ = next(model_.parameters()).device
         except Exception:
             if torch.cuda.is_available():
-                model = model.to(torch.device("cuda"))
-        model.eval()
-        chat_manager.init_if_needed(args, args.model_path, tokenizer, model, image_processor, context_len)
+                model_ = model_.to(torch.device("cuda"))
+        model_.eval()
+
+        # assign globals
+        globals()["tokenizer"] = tokenizer_
+        globals()["model"] = model_
+        globals()["image_processor"] = image_processor_
+        globals()["context_len"] = context_len_
+
+        chat_manager.init_if_needed(args, args.model_path, tokenizer_, model_, image_processor_, context_len_)
         print("[init] model/tokenizer/image_processor loaded.")
         return True
     except Exception as e:
         print(f"[init] failed: {e}")
         return False
 
-# ----------------- HF EndpointHandler -----------------
+
+# ======================== HF EndpointHandler ========================
 
 class EndpointHandler:
     """Hugging Face Endpoint uyumlu sınıf"""