Update handler.py

handler.py

@@ -2,8 +2,8 @@ from typing import Dict, List, Any
 import torch
 import json
 import os
-from transformers import PreTrainedTokenizerFast, PreTrainedModel
-from transformers.models.qwen2 import Qwen2Config, Qwen2ForCausalLM
+import glob
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
 
 # Set up logging
@@ -13,100 +13,135 @@ logger = logging.getLogger(__name__)
 class EndpointHandler:
     def __init__(self, path: str = ""):
         """
-        Initialize handler with manual model loading to bypass auto-detection
+        Initialize handler with robust file discovery
         """
         logger.info(f"Loading model from {path}")
         
         try:
-            # Manual config loading and creation
-            config_path = os.path.join(path, "config.json")
-            
-            if os.path.exists(config_path):
-                with open(config_path, 'r') as f:
-                    config_dict = json.load(f)
-                logger.info(f"Loaded config: {config_dict.get('model_type', 'UNKNOWN')}")
-                
-                # Create Qwen2Config manually
-                config = Qwen2Config(**config_dict)
-            else:
-                logger.warning("No config.json found, using default Qwen2Config")
-                config = Qwen2Config()
-            
-            # Load tokenizer manually without auto-detection
-            logger.info("Loading tokenizer manually...")
-            tokenizer_path = os.path.join(path, "tokenizer.json")
-            
-            if os.path.exists(tokenizer_path):
-                # Load tokenizer from tokenizer.json directly
-                self.tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_path)
-            else:
-                # Try loading from vocab files
-                vocab_path = os.path.join(path, "vocab.json")
-                merges_path = os.path.join(path, "merges.txt")
-                
-                if os.path.exists(vocab_path):
-                    self.tokenizer = PreTrainedTokenizerFast(
-                        tokenizer_file=None,
-                        vocab_file=vocab_path,
-                        merges_file=merges_path if os.path.exists(merges_path) else None
-                    )
-                else:
-                    # Fallback: create basic tokenizer
-                    from transformers import AutoTokenizer
-                    logger.warning("Using fallback tokenizer loading...")
-                    self.tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B-Instruct")
-            
-            # Set special tokens
-            if not hasattr(self.tokenizer, 'pad_token') or self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = "<|endoftext|>"
-                self.tokenizer.pad_token_id = 151643
-            
-            if not hasattr(self.tokenizer, 'eos_token') or self.tokenizer.eos_token is None:
-                self.tokenizer.eos_token = "<|endoftext|>"
-                self.tokenizer.eos_token_id = 151643
+            # Log directory contents to understand structure
+            if os.path.exists(path):
+                contents = os.listdir(path)
+                logger.info(f"Repository contents: {contents}")
                 
+                # Look for model files in subdirectories
+                for item in contents:
+                    item_path = os.path.join(path, item)
+                    if os.path.isdir(item_path):
+                        sub_contents = os.listdir(item_path)
+                        logger.info(f"Directory {item}: {sub_contents}")
+            
+            # Try to find the actual model path
+            model_path = self._find_model_path(path)
+            logger.info(f"Using model path: {model_path}")
+            
+            # Load tokenizer - try multiple approaches
+            self.tokenizer = self._load_tokenizer(model_path, path)
             logger.info("Tokenizer loaded successfully")
             
-            # Load model manually with the config
-            logger.info("Loading model manually...")
-            self.model = Qwen2ForCausalLM(config)
-            
-            # Load state dict manually
-            safetensors_files = [f for f in os.listdir(path) if f.endswith('.safetensors')]
-            
-            if safetensors_files:
-                logger.info(f"Loading weights from {len(safetensors_files)} safetensors files")
-                from safetensors.torch import load_file
-                
-                state_dict = {}
-                for file in sorted(safetensors_files):
-                    file_path = os.path.join(path, file)
-                    partial_state_dict = load_file(file_path)
-                    state_dict.update(partial_state_dict)
-                
-                # Load the state dict
-                missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False)
-                
-                if missing_keys:
-                    logger.warning(f"Missing keys: {missing_keys[:5]}...")  # Show first 5
-                if unexpected_keys:
-                    logger.warning(f"Unexpected keys: {unexpected_keys[:5]}...")  # Show first 5
-            else:
-                logger.error("No safetensors files found!")
-                raise FileNotFoundError("No model weights found")
-            
-            # Move to GPU and set to eval mode
-            self.model = self.model.half()  # Convert to float16
-            if torch.cuda.is_available():
-                self.model = self.model.cuda()
-            self.model.eval()
-            
+            # Load model
+            self.model = self._load_model(model_path, path)
             logger.info("Model loaded successfully")
             
         except Exception as e:
-            logger.error(f"Failed to load model: {str(e)}")
+            logger.error(f"Failed to initialize: {str(e)}")
             raise e
     
+    def _find_model_path(self, base_path: str) -> str:
+        """Find the actual path containing model files"""
+        
+        # Check if config.json is in base path
+        if os.path.exists(os.path.join(base_path, "config.json")):
+            return base_path
+        
+        # Check models/huggingface subdirectory
+        hf_path = os.path.join(base_path, "models", "huggingface")
+        if os.path.exists(hf_path) and os.path.exists(os.path.join(hf_path, "config.json")):
+            return hf_path
+        
+        # Check for any subdirectory with config.json
+        for root, dirs, files in os.walk(base_path):
+            if "config.json" in files:
+                return root
+        
+        # Fallback to base path
+        return base_path
+    
+    def _load_tokenizer(self, model_path: str, base_path: str):
+        """Load tokenizer with fallback methods"""
+        
+        try:
+            # Try direct loading from model path
+            logger.info(f"Trying to load tokenizer from {model_path}")
+            return AutoTokenizer.from_pretrained(
+                model_path,
+                trust_remote_code=True,
+                local_files_only=True
+            )
+        except Exception as e1:
+            logger.warning(f"Failed to load from {model_path}: {e1}")
+            
+            try:
+                # Try loading from base path
+                logger.info(f"Trying to load tokenizer from {base_path}")
+                return AutoTokenizer.from_pretrained(
+                    base_path,
+                    trust_remote_code=True,
+                    local_files_only=True
+                )
+            except Exception as e2:
+                logger.warning(f"Failed to load from {base_path}: {e2}")
+                
+                try:
+                    # Try loading from Hugging Face Hub as fallback
+                    logger.info("Using fallback tokenizer from Qwen2-7B-Instruct")
+                    tokenizer = AutoTokenizer.from_pretrained(
+                        "Qwen/Qwen2-7B-Instruct",
+                        trust_remote_code=True
+                    )
+                    
+                    # Set special tokens
+                    tokenizer.pad_token = tokenizer.eos_token
+                    return tokenizer
+                    
+                except Exception as e3:
+                    logger.error(f"All tokenizer loading methods failed: {e3}")
+                    raise e3
+    
+    def _load_model(self, model_path: str, base_path: str):
+        """Load model with fallback methods"""
+        
+        try:
+            # Try direct loading from model path
+            logger.info(f"Trying to load model from {model_path}")
+            model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True,
+                local_files_only=True,
+                low_cpu_mem_usage=True
+            )
+        except Exception as e1:
+            logger.warning(f"Failed to load from {model_path}: {e1}")
+            
+            try:
+                # Try loading from base path
+                logger.info(f"Trying to load model from {base_path}")
+                model = AutoModelForCausalLM.from_pretrained(
+                    base_path,
+                    torch_dtype=torch.float16,
+                    device_map="auto",
+                    trust_remote_code=True,
+                    local_files_only=True,
+                    low_cpu_mem_usage=True
+                )
+            except Exception as e2:
+                logger.error(f"Model loading failed from both paths: {e2}")
+                raise e2
+        
+        model.eval()
+        return model
+    
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         Handle inference requests
@@ -118,53 +153,75 @@ class EndpointHandler:
             if not inputs:
                 return [{"error": "No input provided", "generated_text": ""}]
             
-            # Generation parameters
+            # Generation parameters with safety limits
             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)
             
-            # Format input
+            # 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"
             
-            # Tokenize
-            input_ids = self.tokenizer.encode(
-                formatted_input,
-                return_tensors="pt",
-                truncation=True,
-                max_length=3072
-            )
-            
-            if torch.cuda.is_available():
-                input_ids = input_ids.cuda()
-            
-            # 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
+            # Tokenize with error handling
+            try:
+                input_ids = self.tokenizer.encode(
+                    formatted_input,
+                    return_tensors="pt",
+                    truncation=True,
+                    max_length=3072
                 )
+            except Exception as e:
+                logger.error(f"Tokenization failed: {e}")
+                return [{"error": f"Tokenization failed: {str(e)}", "generated_text": ""}]
+            
+            if input_ids.size(1) == 0:
+                return [{"error": "Empty input after tokenization", "generated_text": ""}]
+            
+            # Move to model device
+            input_ids = input_ids.to(next(self.model.parameters()).device)
+            
+            # Generate with error handling
+            try:
+                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,
+                        num_return_sequences=1
+                    )
+            except Exception as e:
+                logger.error(f"Generation failed: {e}")
+                return [{"error": f"Generation failed: {str(e)}", "generated_text": ""}]
             
             # Decode response
-            generated_ids = outputs[0][input_ids.size(1):]
-            response = self.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
-            response = response.replace("<|im_end|>", "").strip()
-            
-            return [{
-                "generated_text": response,
-                "generated_tokens": len(generated_ids),
-                "finish_reason": "eos_token" if self.tokenizer.eos_token_id in generated_ids else "length"
-            }]
+            try:
+                generated_ids = outputs[0][input_ids.size(1):]
+                response = self.tokenizer.decode(
+                    generated_ids, 
+                    skip_special_tokens=True
+                ).strip()
+                
+                # Clean up response
+                response = response.replace("<|im_end|>", "").strip()
+                
+                return [{
+                    "generated_text": response,
+                    "generated_tokens": len(generated_ids),
+                    "finish_reason": "eos_token" if self.tokenizer.eos_token_id in generated_ids else "length"
+                }]
+                
+            except Exception as e:
+                logger.error(f"Decoding failed: {e}")
+                return [{"error": f"Decoding failed: {str(e)}", "generated_text": ""}]
             
         except Exception as e:
-            logger.error(f"Generation error: {str(e)}")
-            return [{"error": f"Generation failed: {str(e)}", "generated_text": ""}]
+            logger.error(f"Inference error: {str(e)}")
+            return [{"error": f"Inference failed: {str(e)}", "generated_text": ""}]