Update handler.py

handler.py

@@ -2,7 +2,8 @@ from typing import Dict, List, Any
 import torch
 import json
 import os
-from transformers import Qwen2TokenizerFast, Qwen2ForCausalLM
+from transformers import PreTrainedTokenizerFast, PreTrainedModel
+from transformers.models.qwen2 import Qwen2Config, Qwen2ForCausalLM
 import logging
 
 # Set up logging
@@ -12,87 +13,105 @@ logger = logging.getLogger(__name__)
 class EndpointHandler:
     def __init__(self, path: str = ""):
         """
-        Initialize the handler for Qwen2.5-Coder-7B-Instruct-Omni1.1
-        Explicitly using Qwen2 classes to bypass auto-detection
+        Initialize handler with manual model loading to bypass auto-detection
         """
         logger.info(f"Loading model from {path}")
         
         try:
-            # Check if config exists and log it
+            # 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 = json.load(f)
-                logger.info(f"Found config with model_type: {config.get('model_type', 'MISSING')}")
+                    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 in repository")
-            
-            # Load tokenizer explicitly as Qwen2
-            logger.info("Loading tokenizer as Qwen2TokenizerFast...")
-            self.tokenizer = Qwen2TokenizerFast.from_pretrained(
-                path,
-                trust_remote_code=True,
-                padding_side="left"
-            )
+                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")
             
-            # Ensure proper tokens
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
+            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
                 
             logger.info("Tokenizer loaded successfully")
             
-            # Load model explicitly as Qwen2ForCausalLM
-            logger.info("Loading model as Qwen2ForCausalLM...")
-            self.model = Qwen2ForCausalLM.from_pretrained(
-                path,
-                torch_dtype=torch.float16,
-                device_map="auto",
-                trust_remote_code=True,
-                low_cpu_mem_usage=True
-            )
+            # Load model manually with the config
+            logger.info("Loading model manually...")
+            self.model = Qwen2ForCausalLM(config)
             
-            self.model.eval()
-            logger.info("Model loaded successfully")
+            # Load state dict manually
+            safetensors_files = [f for f in os.listdir(path) if f.endswith('.safetensors')]
             
-        except Exception as e:
-            logger.error(f"Error during initialization: {str(e)}")
-            # Try alternative loading method
-            try:
-                logger.info("Attempting alternative loading method...")
-                
-                # Use the models subdirectory path that we saw in your repo
-                model_path = os.path.join(path, "models", "huggingface") if os.path.exists(os.path.join(path, "models", "huggingface")) else path
-                
-                self.tokenizer = Qwen2TokenizerFast.from_pretrained(
-                    model_path,
-                    trust_remote_code=True,
-                    local_files_only=True
-                )
+            if safetensors_files:
+                logger.info(f"Loading weights from {len(safetensors_files)} safetensors files")
+                from safetensors.torch import load_file
                 
-                if self.tokenizer.pad_token is None:
-                    self.tokenizer.pad_token = self.tokenizer.eos_token
+                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)
                 
-                self.model = Qwen2ForCausalLM.from_pretrained(
-                    model_path,
-                    torch_dtype=torch.float16,
-                    device_map="auto",
-                    trust_remote_code=True,
-                    local_files_only=True
-                )
-                
-                self.model.eval()
-                logger.info("Alternative loading successful")
+                # Load the state dict
+                missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False)
                 
-            except Exception as e2:
-                logger.error(f"Alternative loading also failed: {str(e2)}")
-                raise e2
+                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()
+            
+            logger.info("Model loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Failed to load model: {str(e)}")
+            raise e
     
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         Handle inference requests
         """
         try:
-            # Extract inputs and parameters
             inputs = data.get("inputs", "")
             parameters = data.get("parameters", {})
             
@@ -100,13 +119,12 @@ class EndpointHandler:
                 return [{"error": "No input provided", "generated_text": ""}]
             
             # Generation parameters
-            max_new_tokens = min(parameters.get("max_new_tokens", 512), 1024)  # Cap at 1024
-            temperature = max(0.1, min(parameters.get("temperature", 0.7), 2.0))  # Clamp between 0.1 and 2.0
-            top_p = max(0.1, min(parameters.get("top_p", 0.9), 1.0))  # Clamp between 0.1 and 1.0
+            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)
-            repetition_penalty = max(1.0, min(parameters.get("repetition_penalty", 1.1), 2.0))
             
-            # Format input with Qwen chat template
+            # Format input
             if inputs.startswith("<|im_start|>"):
                 formatted_input = inputs
             else:
@@ -116,15 +134,12 @@ class EndpointHandler:
             input_ids = self.tokenizer.encode(
                 formatted_input,
                 return_tensors="pt",
-                add_special_tokens=False,
                 truncation=True,
-                max_length=3072  # Leave room for generation
+                max_length=3072
             )
             
-            if input_ids.size(1) == 0:
-                return [{"error": "Input tokenization failed", "generated_text": ""}]
-            
-            input_ids = input_ids.to(self.model.device)
+            if torch.cuda.is_available():
+                input_ids = input_ids.cuda()
             
             # Generate
             with torch.no_grad():
@@ -134,23 +149,14 @@ class EndpointHandler:
                     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,
-                    use_cache=True,
-                    num_return_sequences=1
+                    use_cache=True
                 )
             
-            # Decode response (only new tokens)
+            # Decode response
             generated_ids = outputs[0][input_ids.size(1):]
-            response = self.tokenizer.decode(
-                generated_ids,
-                skip_special_tokens=True,
-                clean_up_tokenization_spaces=True
-            )
-            
-            # Clean up response
-            response = response.strip()
+            response = self.tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
             response = response.replace("<|im_end|>", "").strip()
             
             return [{
@@ -161,7 +167,4 @@ class EndpointHandler:
             
         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": ""}]