Update handler.py

handler.py

@@ -3,8 +3,9 @@
 PULSE ECG Handler - Deterministic ECG Analysis Model (app.py uyumlu)
 - Deterministic (do_sample=False, sabit seed)
 - Tek görüntü, LLaVA conv_template + <image> token akışı
-- Model dtype/device ile uyumlu görüntü tensörü
+- Model dtype/device ile uyumlu görüntü tensörü (3D/4D/5D destekli)
 - Sağlam URL/base64 işleme, güvenli logging, opsiyonel HF upload
+- Zorunlu başlık şablonu + min_new_tokens ile tam Step 1–9 çıktısı
 """
 
 import os
@@ -19,7 +20,7 @@ from io import BytesIO
 
 # --- Opsiyonel bağımlılıklar ---
 try:
-    import numpy as np  # isteğe bağlı, kullanılabilir
+    import numpy as np  # isteğe bağlı
 except Exception:
     np = None
 
@@ -55,7 +56,7 @@ except Exception as e:
 
 # Transformers
 try:
-    from transformers import TextIteratorStreamer  # kullanılmıyor ama mevcutsa sorun değil
+    from transformers import TextIteratorStreamer  # mevcutsa sorun değil
     TRANSFORMERS_AVAILABLE = True
 except Exception:
     TRANSFORMERS_AVAILABLE = False
@@ -98,10 +99,23 @@ args = None
 model_initialized = False
 
 # --- Tutarlılık ayarları ---
-# Tutarlılık ayarları
 PROMPT_NORMALIZATION = True
-DEFAULT_ECG_PROMPT = "Perform a detailed ECG interpretation of the provided image. Analyze step by step the rhythm, heart rate, cardiac axis, P waves, PR interval, QRS complex morphology and duration, ST segments, T waves, and QT/QTc interval. Highlight any abnormalities, conduction disturbances, or ischemic changes you detect. Conclude with a structured clinical impression of the overall ECG."
-
+DEFAULT_ECG_PROMPT = (
+    "Perform a detailed ECG interpretation of the provided image. Analyze step by step the rhythm, heart rate, "
+    "cardiac axis, P waves, PR interval, QRS complex morphology and duration, ST segments, T waves, and QT/QTc interval. "
+    "OUTPUT FORMAT (use these exact headings, and include every section even if normal):\n"
+    "Step 1: Rhythm Analysis\n"
+    "Step 2: Heart Rate Analysis\n"
+    "Step 3: Cardiac Axis Analysis\n"
+    "Step 4: P Wave Analysis\n"
+    "Step 5: PR Interval Analysis\n"
+    "Step 6: QRS Complex Analysis\n"
+    "Step 7: ST Segment Analysis\n"
+    "Step 8: T Wave Analysis\n"
+    "Step 9: QT/QTc Interval Analysis\n"
+    "Structured Clinical Impression:\n"
+    "If a section is normal, write 'Normal' and give a brief justification."
+)
 
 # ---------- Yardımcılar ----------
 
@@ -287,22 +301,32 @@ def clear_history():
     except Exception as e:
         return {"error": f"Failed to clear history: {str(e)}"}
 
-# ---------- Cevap üretimi ----------
+# ---------- Prompt inşası ----------
 
 def _build_prompt(chatbot, user_text: str) -> str:
-    # App.py ile aynı: <image> token + kullanıcı metni
-    image_token = DEFAULT_IMAGE_TOKEN
-    inp = image_token + "\n" + user_text
+    # mm_use_im_start_end konfigürasyonuna göre <image> tokenını sarmala
+    try:
+        use_wrap = bool(getattr(chatbot.model.config, "mm_use_im_start_end", False))
+    except Exception:
+        use_wrap = False
+
+    if use_wrap:
+        # <im_start><image></im_end>\n + metin
+        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()
-    return prompt
+    return chatbot.conversation.get_prompt()
 
 def _stop_criteria_from_conv(chatbot, input_ids):
     conv = chatbot.conversation
     stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
     return KeywordsStoppingCriteria([stop_str], chatbot.tokenizer, input_ids)
 
+# ---------- Cevap üretimi ----------
+
 def generate_response(message_text,
                       image_input,
                       max_output_tokens=4096,
@@ -351,45 +375,37 @@ def generate_response(message_text,
     # Model dtype/device
     model_device = next(chatbot.model.parameters()).device
     model_dtype  = next(chatbot.model.parameters()).dtype
-    
+
     # Görüntü tensörü (Tensor/list/tuple + 3D/4D/5D destekli)
     try:
         processed = process_images([image], chatbot.image_processor, chatbot.model.config)
-    
+
         if isinstance(processed, torch.Tensor):
-            # Olası şekiller: (C,H,W), (B,C,H,W), (B,T,C,H,W)
+            # (C,H,W) / (B,C,H,W) / (B,T,C,H,W)
             if processed.ndim == 3:
-                # (C,H,W) -> (1,C,H,W)
-                image_tensor = processed.unsqueeze(0)
+                image_tensor = processed.unsqueeze(0)  # (1,C,H,W)
             elif processed.ndim == 4:
-                # (B,C,H,W)
-                image_tensor = processed
+                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)
+                image_tensor = processed.reshape(b * t, c, h, w)  # (B*T,C,H,W)
             else:
                 return {"error": f"Unexpected image tensor shape: {tuple(processed.shape)}"}
-    
         elif isinstance(processed, (list, tuple)):
             if len(processed) == 0:
                 return {"error": "Image processing returned empty list"}
             first = processed[0]
             if not isinstance(first, torch.Tensor):
                 return {"error": f"Processed image type not tensor: {type(first)}"}
-            # first: (C,H,W) veya (B,C,H,W)
             image_tensor = first.unsqueeze(0) if first.ndim == 3 else first
         else:
             return {"error": f"Unsupported processed type: {type(processed)}"}
-    
-        # Cihaz ve dtype eşle
+
         image_tensor = image_tensor.to(device=model_device, dtype=model_dtype)
-    
+
     except Exception as e:
         return {"error": f"Image processing failed: {str(e)}"}
 
-
-
     # Prompt & tokenizasyon
     prompt = _build_prompt(chatbot, message_text)
     input_ids = tokenizer_image_token(
@@ -405,6 +421,14 @@ def generate_response(message_text,
         torch.cuda.manual_seed(42)
         torch.cuda.manual_seed_all(42)
 
+    # EOS/PAD güvenli al
+    eos_id = chatbot.tokenizer.eos_token_id
+    if eos_id is None:
+        try:
+            eos_id = chatbot.tokenizer.convert_tokens_to_ids("</s>")
+        except Exception:
+            eos_id = 0  # son çare
+
     try:
         with torch.no_grad():
             outputs = chatbot.model.generate(
@@ -412,10 +436,11 @@ def generate_response(message_text,
                 images=image_tensor,
                 do_sample=False,                  # deterministik
                 max_new_tokens=int(max_output_tokens),
+                min_new_tokens=600,               # en az bu kadar üret (step başlıkları garanti)
                 repetition_penalty=float(repetition_penalty),
                 use_cache=False,
-                pad_token_id=chatbot.tokenizer.eos_token_id,
-                eos_token_id=chatbot.tokenizer.eos_token_id,
+                pad_token_id=eos_id,
+                eos_token_id=eos_id,
                 length_penalty=1.0,
                 early_stopping=False,
                 stopping_criteria=[stopping_criteria],
@@ -480,10 +505,13 @@ def query(payload):
             or ""
         )
 
-        # Prompt normalization (ECG + diagnosis içeriyorsa)
+        # Prompt normalization (ECG içeren tüm isteklerde ayrıntılı şablonu zorla)
         if PROMPT_NORMALIZATION and "ecg" in message_text.lower():
             if "concise" in message_text.lower():
-                message_text = "Provide a short, concise clinical summary of the ECG."
+                message_text = (
+                    "Provide a short, concise clinical summary of the ECG. "
+                    "Still cover rhythm, rate, axis, PR, QRS, ST-T, QT/QTc in brief."
+                )
             else:
                 message_text = DEFAULT_ECG_PROMPT
 
@@ -498,7 +526,7 @@ def query(payload):
         # Parametreler
         max_output_tokens = int(payload.get("max_output_tokens",
                                   payload.get("max_new_tokens",
-                                  payload.get("max_tokens", 2048))))
+                                  payload.get("max_tokens", 4096))))
         repetition_penalty = float(payload.get("repetition_penalty", 1.0))
         conv_mode_override = payload.get("conv_mode", None)
 
@@ -576,27 +604,30 @@ def initialize_model():
                 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", "2048"))
+                self.max_new_tokens = int(os.getenv("MAX_NEW_TOKENS", "4096"))
                 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")))
-
-        args = Args()
+        # args globaline ata
+        globals()["args"] = Args()
 
         model_name = get_model_name_from_path(args.model_path)
-        tokenizer, model, image_processor, context_len = load_pretrained_model(
+        loaded = load_pretrained_model(
             args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit
         )
+        globals()["tokenizer"], globals()["model"], globals()["image_processor"], globals()["context_len"] = loaded
 
-        # Device: accelerate devicemap kullanıyorsa ek .to('cuda') gerekmez
+        # Device: accelerate devicemap varsa ek .to('cuda') gerekmeyebilir
         try:
             _ = next(model.parameters()).device
         except Exception:
-            # güvenli taşıma
             if torch.cuda.is_available():
                 model = model.to(torch.device("cuda"))
 
+        # Deterministik için dropout vb. kapansın
+        model.eval()
+
         print("[init] tokenizer/image_processor/context_len ready")
         return True