Update handler.py

handler.py

@@ -1,68 +1,230 @@
-from typing import Dict, List, Any
+from typing import Dict, List, Any, Optional
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import json
+import logging
+
+# Loglama ayarları
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class EndpointHandler:
     def __init__(self, path=""):
         """
-        Qwen2 modelini yükle
+        Qwen2 7.6B modelini optimize edilmiş şekilde yükle
         """
-        # Tokenizer'ı yükle
-        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
-        
-        # Modeli yükle - GPU belleğini optimize et
-        self.model = AutoModelForCausalLM.from_pretrained(
-            path,
-            torch_dtype=torch.bfloat16,  # Bellek optimizasyonu için
-            device_map="auto",            # GPU'yu otomatik kullan
-            trust_remote_code=True        # Qwen2 için gerekli
-        )
-        
-        # Eğer pad token yoksa ekle
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
+        try:
+            logger.info(f"Model yükleniyor: {path}")
+            
+            # Tokenizer'ı yükle - Qwen2 için trust_remote_code gerekli
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                path,
+                trust_remote_code=True,
+                use_fast=True  # Fast tokenizer kullan
+            )
+            
+            # Model konfigürasyonu
+            model_kwargs = {
+                "torch_dtype": torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+                "device_map": "auto",
+                "trust_remote_code": True,
+                "low_cpu_mem_usage": True,  # Bellek optimizasyonu
+            }
+            
+            # Modeli yükle
+            self.model = AutoModelForCausalLM.from_pretrained(
+                path,
+                **model_kwargs
+            )
+            
+            # Model'i eval moduna al
+            self.model.eval()
+            
+            # Tokenizer ayarları
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            if self.tokenizer.pad_token_id is None:
+                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+                
+            # Chat template kontrolü
+            self.has_chat_template = hasattr(self.tokenizer, 'chat_template') and self.tokenizer.chat_template is not None
             
+            logger.info(f"Model başarıyla yüklendi. Chat template: {self.has_chat_template}")
+            logger.info(f"Device: {next(self.model.parameters()).device}")
+            logger.info(f"Dtype: {next(self.model.parameters()).dtype}")
+            
+        except Exception as e:
+            logger.error(f"Model yükleme hatası: {str(e)}")
+            raise RuntimeError(f"Model initialization failed: {str(e)}")
+    
+    def format_chat_input(self, messages: List[Dict[str, str]]) -> str:
+        """
+        Chat formatında gelen mesajları işle
+        """
+        if self.has_chat_template:
+            return self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+        else:
+            # Fallback: Basit format
+            formatted = ""
+            for message in messages:
+                role = message.get("role", "user")
+                content = message.get("content", "")
+                if role == "system":
+                    formatted += f"System: {content}\n"
+                elif role == "user":
+                    formatted += f"User: {content}\n"
+                elif role == "assistant":
+                    formatted += f"Assistant: {content}\n"
+            formatted += "Assistant: "
+            return formatted
+    
+    @torch.inference_mode()
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
-        Inference endpoint
+        Inference endpoint - Hem text hem de chat formatını destekler
         """
-        # Girdileri al
-        inputs = data.pop("inputs", "")
-        
-        # Parametreleri al (varsayılan değerlerle)
-        parameters = data.pop("parameters", {})
-        max_new_tokens = parameters.get("max_new_tokens", 256)
-        temperature = parameters.get("temperature", 0.7)
-        top_p = parameters.get("top_p", 0.95)
-        do_sample = parameters.get("do_sample", True)
-        repetition_penalty = parameters.get("repetition_penalty", 1.1)
-        
-        # Tokenize et
-        input_ids = self.tokenizer(
-            inputs,
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=2048
-        ).input_ids.to(self.model.device)
-        
-        # Generate et
-        with torch.no_grad():
-            outputs = self.model.generate(
-                input_ids,
-                max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                do_sample=do_sample,
-                repetition_penalty=repetition_penalty,
-                pad_token_id=self.tokenizer.pad_token_id,
-                eos_token_id=self.tokenizer.eos_token_id
+        try:
+            # Input türünü belirle
+            inputs = data.pop("inputs", None)
+            messages = data.pop("messages", None)
+            
+            # Input kontrolü
+            if not inputs and not messages:
+                return [{"error": "Either 'inputs' or 'messages' must be provided"}]
+            
+            # Chat format'ı kontrol et
+            if messages:
+                text_input = self.format_chat_input(messages)
+            else:
+                text_input = inputs
+            
+            # Parametreleri al
+            parameters = data.pop("parameters", {})
+            
+            # Generation parametreleri
+            max_new_tokens = parameters.get("max_new_tokens", 256)
+            temperature = parameters.get("temperature", 0.7)
+            top_p = parameters.get("top_p", 0.9)
+            top_k = parameters.get("top_k", 50)
+            do_sample = parameters.get("do_sample", True)
+            repetition_penalty = parameters.get("repetition_penalty", 1.1)
+            num_return_sequences = parameters.get("num_return_sequences", 1)
+            stop_sequences = parameters.get("stop_sequences", None)
+            
+            logger.info(f"Processing input (length: {len(text_input)})")
+            
+            # Tokenize
+            inputs_encoded = self.tokenizer(
+                text_input,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=min(2048, self.model.config.max_position_embeddings),
+                return_attention_mask=True
             )
-        
-        # Sadece yeni oluşturulan tokenleri al
-        generated_ids = outputs[0][input_ids.shape[-1]:]
-        
-        # Decode et
-        result = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
-        
-        return [{"generated_text": result}]
+            
+            # Device'a taşı
+            input_ids = inputs_encoded["input_ids"].to(self.model.device)
+            attention_mask = inputs_encoded["attention_mask"].to(self.model.device)
+            
+            # Stopping criteria ayarla
+            stop_token_ids = []
+            if stop_sequences:
+                for seq in stop_sequences:
+                    tokens = self.tokenizer.encode(seq, add_special_tokens=False)
+                    stop_token_ids.extend(tokens)
+            
+            # Generate
+            generation_kwargs = {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "max_new_tokens": max_new_tokens,
+                "temperature": temperature if do_sample else 1.0,
+                "top_p": top_p if do_sample else 1.0,
+                "top_k": top_k if do_sample else None,
+                "do_sample": do_sample,
+                "repetition_penalty": repetition_penalty,
+                "num_return_sequences": num_return_sequences,
+                "pad_token_id": self.tokenizer.pad_token_id,
+                "eos_token_id": self.tokenizer.eos_token_id,
+                "use_cache": True,
+            }
+            
+            # Stop tokens ekle
+            if stop_token_ids:
+                generation_kwargs["eos_token_id"] = stop_token_ids
+            
+            # Generate
+            outputs = self.model.generate(**generation_kwargs)
+            
+            # Decode
+            results = []
+            for output in outputs:
+                # Input kısmını çıkar
+                generated_ids = output[input_ids.shape[-1]:]
+                generated_text = self.tokenizer.decode(
+                    generated_ids,
+                    skip_special_tokens=True,
+                    clean_up_tokenization_spaces=True
+                )
+                
+                results.append({
+                    "generated_text": generated_text,
+                    "details": {
+                        "finish_reason": "length" if len(generated_ids) >= max_new_tokens else "stop",
+                        "generated_tokens": len(generated_ids),
+                        "input_tokens": input_ids.shape[-1]
+                    }
+                })
+            
+            logger.info(f"Generation completed. Generated {len(results)} sequences")
+            
+            # Tek sonuç istenmişse direkt döndür
+            if num_return_sequences == 1:
+                return results
+            else:
+                return [{"results": results}]
+            
+        except torch.cuda.OutOfMemoryError:
+            logger.error("GPU bellek yetersiz!")
+            return [{
+                "error": "GPU out of memory. Try reducing max_new_tokens or input length",
+                "type": "memory_error"
+            }]
+        except Exception as e:
+            logger.error(f"Inference hatası: {str(e)}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return [{
+                "error": str(e),
+                "type": "inference_error",
+                "traceback": traceback.format_exc()
+            }]
+    
+    def health_check(self) -> Dict[str, Any]:
+        """
+        Endpoint sağlık kontrolü
+        """
+        try:
+            test_input = "Test"
+            inputs = self.tokenizer(test_input, return_tensors="pt")
+            with torch.no_grad():
+                _ = self.model.generate(
+                    inputs.input_ids.to(self.model.device),
+                    max_new_tokens=5
+                )
+            return {
+                "status": "healthy",
+                "model": "Qwen2-7.6B",
+                "device": str(next(self.model.parameters()).device),
+                "dtype": str(next(self.model.parameters()).dtype)
+            }
+        except Exception as e:
+            return {
+                "status": "unhealthy",
+                "error": str(e)
+            }