Fix: Add pad_token and attention_mask for proper generation"

3612e19 verified 11 months ago

3.28 kB

from typing import Dict, Any
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


class EndpointHandler:
    def __init__(self, path: str = ""):
        """
        Initialize the model and tokenizer.
        
        Args:
            path: Path to the model directory (will be "/repository" in endpoint container)
        """
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        print(f"Loading tokenizer from {path}...")
        self.tokenizer = AutoTokenizer.from_pretrained(path)

        # ✅ ADD THIS: Set pad token to prevent corruption
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        
        print(f"Loading model from {path} on device: {self.device}...")
        self.model = AutoModelForCausalLM.from_pretrained(
            path,
            torch_dtype=torch.float16,  # Your merged model is fp16
            trust_remote_code=True,     # StarCoder2 may use custom code
            device_map="auto",          # Efficient placement on GPU/CPU
        )
        
        self.model.eval()  # Set to evaluation mode
        print("✅ Model loaded successfully!")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process inference requests.
        
        Args:
            data: Dictionary containing:
                - inputs: str (code prompt to complete)
                - parameters: dict (optional, generation parameters)
        
        Returns:
            Dictionary with generated_text key
        """
        # Extract inputs
        inputs = data.get("inputs", "")
        parameters = data.get("parameters", {}) or {}
        
        if not isinstance(inputs, str):
            raise ValueError("`inputs` must be a string")
        
        if not inputs.strip():
            raise ValueError("`inputs` cannot be empty")
        
        # Generation parameters with sensible defaults
        gen_kwargs = {
            "max_new_tokens": parameters.get("max_new_tokens", 128),
            "temperature": parameters.get("temperature", 0.2),  # Lower = more deterministic
            "top_p": parameters.get("top_p", 0.95),             # Nucleus sampling
            "top_k": parameters.get("top_k", 50),               # Top-k sampling
            "do_sample": parameters.get("do_sample", True),     # Use sampling
            "repetition_penalty": parameters.get("repetition_penalty", 1.0),
        }
        
        print(f"Generating with parameters: {gen_kwargs}")
        
        # Tokenize input
        enc = self.tokenizer(inputs, return_tensors="pt",padding=True,           # Enable padding
        truncation=True,        #  Truncate if needed
        max_length=2048).to(self.device)
        
        # Generate with no_grad for efficiency
        with torch.no_grad():
            out = self.model.generate(**enc, 
                                      **gen_kwargs,
                                      pad_token_id=self.tokenizer.pad_token_id) # Tell model which token is padding
        
        # Decode output
        generated_text = self.tokenizer.decode(out[0], skip_special_tokens=True)
        
        return {
            "generated_text": generated_text
        }

	from typing import Dict, Any
	import torch
	from transformers import AutoModelForCausalLM, AutoTokenizer


	class EndpointHandler:
	def __init__(self, path: str = ""):
	"""
	Initialize the model and tokenizer.

	Args:
	path: Path to the model directory (will be "/repository" in endpoint container)
	"""
	self.device = "cuda" if torch.cuda.is_available() else "cpu"

	print(f"Loading tokenizer from {path}...")
	self.tokenizer = AutoTokenizer.from_pretrained(path)

	# ✅ ADD THIS: Set pad token to prevent corruption
	if self.tokenizer.pad_token is None:
	self.tokenizer.pad_token = self.tokenizer.eos_token

	print(f"Loading model from {path} on device: {self.device}...")
	self.model = AutoModelForCausalLM.from_pretrained(
	path,
	torch_dtype=torch.float16, # Your merged model is fp16
	trust_remote_code=True, # StarCoder2 may use custom code
	device_map="auto", # Efficient placement on GPU/CPU
	)

	self.model.eval() # Set to evaluation mode
	print("✅ Model loaded successfully!")

	def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
	"""
	Process inference requests.

	Args:
	data: Dictionary containing:
	- inputs: str (code prompt to complete)
	- parameters: dict (optional, generation parameters)

	Returns:
	Dictionary with generated_text key
	"""
	# Extract inputs
	inputs = data.get("inputs", "")
	parameters = data.get("parameters", {}) or {}

	if not isinstance(inputs, str):
	raise ValueError("`inputs` must be a string")

	if not inputs.strip():
	raise ValueError("`inputs` cannot be empty")

	# Generation parameters with sensible defaults
	gen_kwargs = {
	"max_new_tokens": parameters.get("max_new_tokens", 128),
	"temperature": parameters.get("temperature", 0.2), # Lower = more deterministic
	"top_p": parameters.get("top_p", 0.95), # Nucleus sampling
	"top_k": parameters.get("top_k", 50), # Top-k sampling
	"do_sample": parameters.get("do_sample", True), # Use sampling
	"repetition_penalty": parameters.get("repetition_penalty", 1.0),
	}

	print(f"Generating with parameters: {gen_kwargs}")

	# Tokenize input
	enc = self.tokenizer(inputs, return_tensors="pt",padding=True, # Enable padding
	truncation=True, # Truncate if needed
	max_length=2048).to(self.device)

	# Generate with no_grad for efficiency
	with torch.no_grad():
	out = self.model.generate(**enc,
	**gen_kwargs,
	pad_token_id=self.tokenizer.pad_token_id) # Tell model which token is padding

	# Decode output
	generated_text = self.tokenizer.decode(out[0], skip_special_tokens=True)

	return {
	"generated_text": generated_text
	}