Update handler.py

handler.py

@@ -11,23 +11,24 @@ logger = logging.getLogger(__name__)
 class EndpointHandler:
     def __init__(self, path: str = ""):
         """
-        Initialize handler that completely bypasses HF auto-detection
+        Initialize handler using CTransformers format for memory efficiency
         """
         logger.info(f"Loading model from {path}")
         
         try:
-            # Set environment variables to avoid auto-detection issues
-            os.environ['TRANSFORMERS_OFFLINE'] = '1'
-            os.environ['HF_DATASETS_OFFLINE'] = '1'
-            os.environ['HF_HUB_OFFLINE'] = '1'
+            # Use CTransformers format for lower memory usage
+            ctransformers_path = os.path.join(path, "models", "ctransformers")
             
-            # Find model files
-            model_path = self._discover_model_files(path)
-            logger.info(f"Model files found at: {model_path}")
+            if not os.path.exists(ctransformers_path):
+                logger.warning(f"CTransformers path not found: {ctransformers_path}")
+                logger.info("Falling back to HuggingFace format")
+                ctransformers_path = path
             
-            # Load components manually
-            self.tokenizer = self._load_tokenizer_manual(model_path)
-            self.model = self._load_model_manual(model_path)
+            logger.info(f"Using model path: {ctransformers_path}")
+            
+            # Load components using the working handler approach
+            self.tokenizer = self._load_tokenizer(ctransformers_path)
+            self.model = self._load_model(ctransformers_path)
             
             logger.info("Model and tokenizer loaded successfully")
             
@@ -35,239 +36,236 @@ class EndpointHandler:
             logger.error(f"Failed to initialize: {str(e)}")
             raise e
     
-    def _discover_model_files(self, base_path: str) -> str:
-        """Find where the actual model files are located"""
-        
-        logger.info(f"Searching for model files in: {base_path}")
-        
-        # List all contents
-        if os.path.exists(base_path):
-            contents = os.listdir(base_path)
-            logger.info(f"Base directory contents: {contents}")
-            
-            # Check for config.json in base path
-            if "config.json" in contents:
-                logger.info("Found config.json in base directory")
-                return base_path
-            
-            # Check models subdirectories
-            for item in contents:
-                if os.path.isdir(os.path.join(base_path, item)):
-                    sub_path = os.path.join(base_path, item)
-                    sub_contents = os.listdir(sub_path)
-                    logger.info(f"Subdirectory {item}: {sub_contents}")
-                    
-                    if "config.json" in sub_contents:
-                        logger.info(f"Found config.json in {item} subdirectory")
-                        return sub_path
-            
-            # Search recursively
-            for root, dirs, files in os.walk(base_path):
-                if "config.json" in files:
-                    logger.info(f"Found config.json in {root}")
-                    return root
-        
-        raise FileNotFoundError(f"No config.json found in {base_path} or subdirectories")
-    
-    def _load_tokenizer_manual(self, model_path: str):
-        """Load tokenizer completely manually"""
+    def _load_tokenizer(self, model_path: str):
+        """Load tokenizer using AutoTokenizer"""
+        logger.info("Loading tokenizer...")
         
-        logger.info("Loading tokenizer manually...")
+        from transformers import AutoTokenizer
         
-        # Check what tokenizer files exist
-        tokenizer_files = []
-        for file in os.listdir(model_path):
-            if file in ['tokenizer.json', 'tokenizer_config.json', 'vocab.json']:
-                tokenizer_files.append(file)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            trust_remote_code=True,
+            use_fast=True,
+        )
         
-        logger.info(f"Found tokenizer files: {tokenizer_files}")
-        
-        if 'tokenizer.json' in tokenizer_files:
-            # Load from tokenizer.json directly
-            from transformers import PreTrainedTokenizerFast
-            tokenizer_path = os.path.join(model_path, 'tokenizer.json')
-            logger.info(f"Loading tokenizer from {tokenizer_path}")
-            
-            tokenizer = PreTrainedTokenizerFast(
-                tokenizer_file=tokenizer_path,
-                unk_token="<|endoftext|>",
-                bos_token="<|endoftext|>",
-                eos_token="<|endoftext|>"
-            )
-        else:
-            # Create a basic tokenizer
-            logger.warning("No tokenizer.json found, creating basic tokenizer")
-            from transformers import PreTrainedTokenizerFast
-            
-            # Create minimal tokenizer
-            tokenizer = PreTrainedTokenizerFast(
-                tokenizer_file=None,
-                vocab_size=151936,  # Qwen2 default vocab size
-                unk_token="<|endoftext|>",
-                bos_token="<|endoftext|>",
-                eos_token="<|endoftext|>",
-                pad_token="<|endoftext|>"
-            )
-        
-        # Set special tokens
+        # Ensure special tokens are set
         if not hasattr(tokenizer, 'pad_token') or tokenizer.pad_token is None:
-            tokenizer.pad_token = "<|endoftext|>"
-            tokenizer.pad_token_id = 151643
-        
-        if not hasattr(tokenizer, 'eos_token') or tokenizer.eos_token is None:
-            tokenizer.eos_token = "<|endoftext|>"
-            tokenizer.eos_token_id = 151643
+            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.pad_token_id = tokenizer.eos_token_id
         
+        logger.info("Tokenizer loaded successfully")
         return tokenizer
     
-    def _load_model_manual(self, model_path: str):
-        """Load model completely manually with memory optimization"""
+    def _load_model(self, model_path: str):
+        """Load model using AutoModelForCausalLM with memory optimization"""
+        logger.info("Loading model with memory optimization...")
         
-        logger.info("Loading model manually...")
+        from transformers import AutoModelForCausalLM, BitsAndBytesConfig
         
-        # Check GPU availability and memory
+        # Check GPU availability
         if torch.cuda.is_available():
             logger.info(f"CUDA available: {torch.cuda.get_device_name()}")
-            logger.info(f"GPU memory before loading: {torch.cuda.memory_allocated() / 1024**3:.2f}GB")
             logger.info(f"GPU memory total: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f}GB")
         else:
             logger.warning("CUDA not available, using CPU")
         
-        # Load config manually
-        config_path = os.path.join(model_path, "config.json")
-        with open(config_path, 'r') as f:
-            config_dict = json.load(f)
+        # Memory optimization settings
+        device_map = "auto" if torch.cuda.is_available() else None
+        gpu_mem = os.environ.get("GPU_MAX_MEM", "10GiB")  # Conservative for 12GB limit
+        cpu_mem = os.environ.get("CPU_MAX_MEM", "24GiB")
+        max_memory = {0: gpu_mem, "cpu": cpu_mem} if torch.cuda.is_available() else None
+        
+        # Offload folder for memory management
+        offload_folder = os.environ.get("OFFLOAD_FOLDER", "/tmp/hf-offload")
+        try:
+            os.makedirs(offload_folder, exist_ok=True)
+        except OSError:
+            offload_folder = "/tmp/hf-offload"
+            os.makedirs(offload_folder, exist_ok=True)
         
-        logger.info(f"Config loaded: {config_dict.get('model_type', 'UNKNOWN')}")
+        # 8-bit quantization for memory efficiency
+        bnb_config = BitsAndBytesConfig(load_in_8bit=True)
         
-        # Create model architecture manually
-        from transformers.models.qwen2 import Qwen2Config, Qwen2ForCausalLM
+        # Load model with all optimizations
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            trust_remote_code=True,
+            device_map=device_map,
+            quantization_config=bnb_config,
+            low_cpu_mem_usage=True,
+            offload_folder=offload_folder if device_map == "auto" else None,
+            max_memory=max_memory,
+        )
         
-        # Ensure model_type is set correctly
-        if 'model_type' not in config_dict:
-            config_dict['model_type'] = 'qwen2'
-            logger.info("Set model_type to 'qwen2'")
+        model.eval()
         
-        # Create config object
-        config = Qwen2Config(**config_dict)
+        # Set context window
+        self.max_context = getattr(model.config, "max_position_embeddings", None) or getattr(self.tokenizer, "model_max_length", 4096)
+        if self.max_context is None or self.max_context == int(1e30):
+            self.max_context = 4096
         
-        # Create model
-        model = Qwen2ForCausalLM(config)
-        logger.info("Model architecture created")
+        # Set token IDs
+        self.pad_token_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id
+        self.eos_token_id = self.tokenizer.eos_token_id
+        
+        logger.info("Model loaded successfully with memory optimization")
+        return model
+    
+    def _build_prompt(self, data: Dict[str, Any]) -> str:
+        """Build prompt using chat template or direct input"""
+        # Accept either "messages" (chat) or "inputs"/"prompt" (single-turn)
+        if "messages" in data and isinstance(data["messages"], list):
+            return self.tokenizer.apply_chat_template(
+                data["messages"],
+                tokenize=False,
+                add_generation_prompt=True
+            )
+        
+        user_text = data.get("inputs") or data.get("prompt") or ""
+        if isinstance(user_text, str):
+            messages = [{"role": "user", "content": user_text}]
+            return self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+        
+        return str(user_text)
+    
+    def _prepare_inputs(self, prompt: str, max_new_tokens: int, params: Dict[str, Any]) -> Dict[str, torch.Tensor]:
+        """Prepare inputs with proper tokenization"""
+        # Keep room for generation
+        max_input_tokens = int(params.get("max_input_tokens", max(self.max_context - max_new_tokens - 8, 256)))
+        
+        model_inputs = self.tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=max_input_tokens,
+        )
         
         if torch.cuda.is_available():
-            logger.info(f"GPU memory after model creation: {torch.cuda.memory_allocated() / 1024**3:.2f}GB")
+            model_inputs = {k: v.to(self.model.device) for k, v in model_inputs.items()}
         
-        # Load weights manually from safetensors with memory optimization
-        import glob
-        safetensors_files = glob.glob(os.path.join(model_path, "*.safetensors"))
-        logger.info(f"Found {len(safetensors_files)} safetensors files")
+        return model_inputs
+    
+    def _stopping(self, params: Dict[str, Any]):
+        """Create stopping criteria"""
+        from transformers import StoppingCriteria, StoppingCriteriaList
         
-        if safetensors_files:
-            from safetensors.torch import load_file
+        class StopOnSequences(StoppingCriteria):
+            def __init__(self, stop_sequences: List[List[int]]):
+                super().__init__()
+                self.stop_sequences = [torch.tensor(x, dtype=torch.long) for x in stop_sequences if len(x) > 0]
             
-            # Load weights directly into model without accumulating in state_dict
-            for i, file in enumerate(sorted(safetensors_files)):
-                logger.info(f"Loading weights from file {i+1}/{len(safetensors_files)}: {os.path.basename(file)}")
-                
-                # Load partial weights
-                partial_state_dict = load_file(file)
-                
-                if torch.cuda.is_available():
-                    logger.info(f"GPU memory after loading file {i+1}: {torch.cuda.memory_allocated() / 1024**3:.2f}GB")
-                
-                # Load this partial state dict directly into the model
-                missing_keys, unexpected_keys = model.load_state_dict(partial_state_dict, strict=False)
-                
-                # Clear partial dict immediately to free memory
-                del partial_state_dict
-                
-                # Force garbage collection
-                import gc
-                gc.collect()
-                
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
-                    logger.info(f"GPU memory after cleanup: {torch.cuda.memory_allocated() / 1024**3:.2f}GB")
+            def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+                if input_ids.shape[0] == 0 or not self.stop_sequences:
+                    return False
+                generated = input_ids[0]
+                for seq in self.stop_sequences:
+                    if generated.shape[0] >= seq.shape[0] and torch.equal(generated[-seq.shape[0]:], seq.to(generated.device)):
+                        return True
+                return False
         
-        # Convert to half precision and move to GPU
-        logger.info("Converting model to half precision...")
-        model = model.half()
+        stop = params.get("stop", [])
+        if isinstance(stop, str):
+            stop = [stop]
+        if not isinstance(stop, list):
+            stop = []
         
-        if torch.cuda.is_available():
-            logger.info(f"GPU memory after half precision: {torch.cuda.memory_allocated() / 1024**3:.2f}GB")
-            model = model.cuda()
-            logger.info(f"GPU memory after moving to GPU: {torch.cuda.memory_allocated() / 1024**3:.2f}GB")
+        stop_ids = [self.tokenizer.encode(s, add_special_tokens=False) for s in stop]
+        criteria = []
+        if stop_ids:
+            criteria.append(StopOnSequences(stop_ids))
         
-        model.eval()
-        logger.info("Model loaded successfully and set to eval mode")
-        return model
+        return StoppingCriteriaList(criteria)
     
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        """
-        Handle inference requests
-        """
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Handle inference requests with proper error handling"""
         try:
-            inputs = data.get("inputs", "")
-            parameters = data.get("parameters", {})
+            params = data.get("parameters", {}) or {}
             
-            if not inputs:
-                return [{"error": "No input provided", "generated_text": ""}]
+            # Set seed if provided
+            seed = params.get("seed")
+            if seed is not None:
+                try:
+                    torch.manual_seed(int(seed))
+                except (ValueError, TypeError):
+                    pass
             
             # Generation parameters
-            max_new_tokens = min(parameters.get("max_new_tokens", 512), 1024)
-            temperature = max(0.1, min(parameters.get("temperature", 0.7), 2.0))
-            top_p = max(0.1, min(parameters.get("top_p", 0.9), 1.0))
-            do_sample = parameters.get("do_sample", True)
+            max_new_tokens = int(params.get("max_new_tokens", 512))
+            temperature = float(params.get("temperature", 0.2))
+            top_p = float(params.get("top_p", 0.9))
+            top_k = int(params.get("top_k", 50))
+            repetition_penalty = float(params.get("repetition_penalty", 1.05))
+            num_beams = int(params.get("num_beams", 1))
+            do_sample = bool(params.get("do_sample", temperature > 0 and num_beams == 1))
             
-            # Format input for Qwen chat template
-            if inputs.startswith("<|im_start|>"):
-                formatted_input = inputs
-            else:
-                formatted_input = f"<|im_start|>user\n{inputs}<|im_end|>\n<|im_start|>assistant\n"
+            # Build prompt
+            prompt = self._build_prompt(data)
+            model_inputs = self._prepare_inputs(prompt, max_new_tokens, params)
+            input_length = model_inputs["input_ids"].shape[-1]
             
-            # Tokenize
-            input_ids = self.tokenizer.encode(
-                formatted_input,
-                return_tensors="pt",
-                truncation=True,
-                max_length=3072
+            # Generation kwargs
+            gen_kwargs = dict(
+                max_new_tokens=max_new_tokens,
+                do_sample=do_sample,
+                temperature=max(0.0, temperature),
+                top_p=top_p,
+                top_k=top_k,
+                repetition_penalty=repetition_penalty,
+                num_beams=num_beams,
+                eos_token_id=self.eos_token_id,
+                pad_token_id=self.pad_token_id,
+                stopping_criteria=self._stopping(params),
             )
             
-            if input_ids.size(1) == 0:
-                return [{"error": "Empty input after tokenization", "generated_text": ""}]
-            
-            input_ids = input_ids.to(self.model.device)
-            
             # Generate
             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,
-                    pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id,
-                    use_cache=True
-                )
+                output_ids = self.model.generate(**model_inputs, **gen_kwargs)
+            
+            # Slice off the prompt
+            gen_ids = output_ids[0][input_length:]
+            text = self.tokenizer.decode(gen_ids, skip_special_tokens=True)
             
-            # Decode response
-            generated_ids = outputs[0][input_ids.size(1):]
-            response = self.tokenizer.decode(
-                generated_ids, 
-                skip_special_tokens=True
-            ).strip()
+            # Apply text-side stop strings if provided
+            stop = params.get("stop", [])
+            if isinstance(stop, str):
+                stop = [stop]
+            for s in stop or []:
+                idx = text.find(s)
+                if idx != -1:
+                    text = text[:idx]
+                    break
             
-            # Clean up response
-            response = response.replace("<|im_end|>", "").strip()
+            # Token accounting
+            prompt_tokens = int(input_length)
+            completion_tokens = int(gen_ids.shape[-1])
+            total_tokens = prompt_tokens + completion_tokens
             
-            return [{
-                "generated_text": response,
-                "generated_tokens": len(generated_ids),
-                "finish_reason": "eos_token" if self.tokenizer.eos_token_id in generated_ids else "length"
-            }]
+            return {
+                "generated_text": text,
+                "input_tokens": prompt_tokens,
+                "generated_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+                "params": {
+                    "max_new_tokens": max_new_tokens,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "top_k": top_k,
+                    "repetition_penalty": repetition_penalty,
+                    "num_beams": num_beams,
+                    "do_sample": do_sample,
+                },
+            }
             
         except Exception as e:
             logger.error(f"Generation error: {str(e)}")
-            return [{"error": f"Generation failed: {str(e)}", "generated_text": ""}]
+            return {
+                "error": f"Generation failed: {str(e)}", 
+                "generated_text": "",
+                "input_tokens": 0,
+                "generated_tokens": 0,
+                "total_tokens": 0
+            }