add custom handler

handler.py

@@ -0,0 +1,152 @@
+from typing import Dict, List, Any, Union
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+import json
+import logging
+
+logger = logging.getLogger(__name__)
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        """
+        Initialize the handler with the PEFT adapter model
+        """
+        try:
+            # Load adapter config to get base model info
+            with open(f"{path}/adapter_config.json", "r") as f:
+                adapter_config = json.load(f)
+            
+            base_model_name = adapter_config.get("base_model_name_or_path", "meta-llama/Llama-2-7b-chat-hf")
+            logger.info(f"Loading base model: {base_model_name}")
+            
+            # Load tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            # Load base model
+            base_model = AutoModelForCausalLM.from_pretrained(
+                base_model_name,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                trust_remote_code=True
+            )
+            
+            # Load PEFT adapter
+            self.model = PeftModel.from_pretrained(
+                base_model, 
+                path,
+                torch_dtype=torch.bfloat16,
+                device_map="auto"
+            )
+            
+            self.model.eval()
+            logger.info("Model loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Error loading model: {e}")
+            raise
+    
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Handle inference requests
+        
+        Expected input format:
+        {
+            "inputs": "build a red car",
+            "parameters": {
+                "max_new_tokens": 512,
+                "temperature": 0.7,
+                "do_sample": true
+            }
+        }
+        """
+        try:
+            # Extract inputs
+            inputs = data.pop("inputs", "")
+            if isinstance(inputs, list):
+                inputs = inputs[0] if inputs else ""
+            
+            # Extract parameters
+            parameters = data.pop("parameters", {})
+            max_new_tokens = parameters.get("max_new_tokens", 512)
+            temperature = parameters.get("temperature", 0.7)
+            do_sample = parameters.get("do_sample", True)
+            top_p = parameters.get("top_p", 0.9)
+            top_k = parameters.get("top_k", 50)
+            
+            # Format prompt for BrickGPT (based on the training format)
+            formatted_prompt = self._format_prompt(inputs)
+            
+            # Tokenize
+            input_ids = self.tokenizer.encode(
+                formatted_prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=2048
+            ).to(self.model.device)
+            
+            # Generate
+            with torch.no_grad():
+                output_ids = self.model.generate(
+                    input_ids,
+                    max_new_tokens=max_new_tokens,
+                    temperature=temperature,
+                    do_sample=do_sample,
+                    top_p=top_p,
+                    top_k=top_k,
+                    pad_token_id=self.tokenizer.eos_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    repetition_penalty=1.1
+                )
+            
+            # Decode only the generated part
+            generated_ids = output_ids[0][input_ids.shape[1]:]
+            generated_text = self.tokenizer.decode(
+                generated_ids,
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=True
+            )
+            
+            # Parse the generated LDR content
+            ldr_instructions = self._parse_ldr_output(generated_text)
+            
+            return [{
+                "generated_text": generated_text,
+                "ldr_instructions": ldr_instructions
+            }]
+            
+        except Exception as e:
+            logger.error(f"Error during inference: {e}")
+            return [{"error": str(e)}]
+    
+    def _format_prompt(self, user_input: str) -> str:
+        """
+        Format the input prompt for BrickGPT
+        Based on how the model was trained
+        """
+        # Clean and format the input
+        if not user_input:
+            user_input = "build something creative"
+        
+        # Format similar to training data
+        prompt = f"### Instruction:\nGenerate LEGO building instructions for: {user_input}\n\n### Response:\n"
+        return prompt
+    
+    def _parse_ldr_output(self, generated_text: str) -> List[str]:
+        """
+        Parse LDR instructions from generated text
+        """
+        ldr_lines = []
+        lines = generated_text.strip().split('\n')
+        
+        for line in lines:
+            line = line.strip()
+            # LDR format lines typically start with numbers or specific commands
+            if line and (line.startswith('1 ') or line.startswith('0 ') or 
+                        line.startswith('2 ') or line.startswith('3 ') or 
+                        line.startswith('4 ') or line.startswith('5 ')):
+                ldr_lines.append(line)
+        
+        return ldr_lines

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.0.0
+transformers>=4.35.0
+peft>=0.6.0
+accelerate>=0.20.0