fix me up dady

handler.py

@@ -1,9 +1,6 @@
 from typing import Dict, List, Any
 import torch
-import re
 import os
-import json
-from pathlib import Path
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
@@ -27,108 +24,37 @@ class EndpointHandler:
         # Set pad token if not exists
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
-        
-        # Load few-shot examples (same as in BrickGPT)
-        self.few_shot_examples = [
-            {
-                "caption": "Bed with rectangular base and straight headboard.",
-                "bricks": "1x2 (13,18,0)\n1x2 (13,2,0)\n2x2 (0,18,0)\n2x2 (0,2,0)\n2x6 (12,14,1)\n2x6 (12,8,1)"
-            },
-            {
-                "caption": "Simple chair with straight backrest and square seat.",
-                "bricks": "2x2 (5,18,0)\n2x2 (5,13,0)\n2x2 (0,18,0)\n2x2 (0,13,0)\n2x2 (5,18,1)\n2x2 (5,13,1)"
-            },
-            {
-                "caption": "Square table with four legs and a flat surface.",
-                "bricks": "2x2 (16,18,0)\n2x2 (16,8,0)\n1x1 (15,18,0)\n1x1 (15,9,0)\n2x2 (0,18,0)\n2x2 (0,8,0)"
-            }
-        ]
-
-    def create_instruction(self, caption: str) -> str:
-        """Create instruction exactly like BrickGPT does"""
-        instruction = ('Create a LEGO model of the input. Format your response as a list of bricks: '
-                       '<brick dimensions> <brick position>, where the brick position is (x,y,z).\n'
-                       'Allowed brick dimensions are 2x4, 4x2, 2x6, 6x2, 1x2, 2x1, 1x4, 4x1, 1x6, 6x1, 1x8, 8x1, 1x1, 2x2.\n'
-                       'All bricks are 1 unit tall.\n\n'
-                       '### Input:\n'
-                       f'{caption}')
-        return instruction
-
-    def create_instruction_few_shot(self, caption: str) -> str:
-        """Create few-shot instruction exactly like BrickGPT does"""
-        base_instruction = self.create_instruction(caption)
-        zero_shot_instructions = (
-            'Each line of your output should be a LEGO brick in the format `<brick dimensions> <brick position>`. For example:\n'
-            '2x4 (2,1,0)\n'
-            'DO NOT output any other text. Only output LEGO bricks. The first brick should have a z-coordinate of 0.'
-        )
-        
-        example_prompt = 'Here are some example LEGO models:'
-        example_instructions = '\n\n'.join(self._create_example_instruction(x) for x in self.few_shot_examples)
-        few_shot_instructions = (
-            'Do NOT copy the examples, but create your own LEGO model for the following input.\n\n'
-            '### Input:\n'
-            f'{caption}\n\n'
-            '### Output:\n'
-        )
-        
-        return '\n\n'.join([base_instruction, zero_shot_instructions, example_prompt,
-                            example_instructions, few_shot_instructions])
-
-    def _create_example_instruction(self, x: dict) -> str:
-        return f'### Input:\n{x["caption"]}\n\n### Output:\n{x["bricks"]}'
-
-    def extract_lego_instructions(self, text: str) -> List[str]:
-        """Extract LEGO brick instructions from generated text"""
-        instructions = []
-        lines = text.split('\n')
-        
-        for line in lines:
-            line = line.strip()
-            if not line:
-                continue
-                
-            # Look for BrickGPT format: "NxM (x,y,z)"
-            brick_pattern = r'(\d+x\d+)\s*\((\d+),(\d+),(\d+)\)'
-            match = re.search(brick_pattern, line)
-            if match:
-                instructions.append(line)
-                
-        return instructions
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
-        Process inference request EXACTLY like local BrickGPT does
+        Simple handler that mimics local LLM behavior for RemoteLLM
         """
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
         
-        # Handle different input formats that BrickGPT sends
+        # Handle different input formats that RemoteLLM sends
         if isinstance(inputs, dict) and "messages" in inputs:
             messages = inputs["messages"]
         elif isinstance(inputs, list):
             messages = inputs
-        elif isinstance(inputs, str):
+        else:
+            # Fallback - treat as direct text
             messages = [
                 {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": inputs}
+                {"role": "user", "content": str(inputs)}
             ]
-        else:
-            messages = [{"role": "user", "content": str(inputs)}]
         
         # Check if this is a continuation (has assistant message)
         has_assistant = any(msg.get("role") == "assistant" for msg in messages)
         
-        # Format prompt EXACTLY like local BrickGPT does
+        # Apply chat template exactly like BrickGPT does locally
         if has_assistant:
-            # For continuation, use continue_final_message=True and return tensors
             prompt = self.tokenizer.apply_chat_template(
                 messages, 
                 continue_final_message=True, 
                 return_tensors='pt'
             )
         else:
-            # For new generation, add generation prompt and return tensors
             prompt = self.tokenizer.apply_chat_template(
                 messages, 
                 add_generation_prompt=True, 
@@ -139,20 +65,19 @@ class EndpointHandler:
         input_ids = prompt.to(self.model.device)
         attention_mask = torch.ones_like(input_ids)
         
-        # Generate EXACTLY like local BrickGPT's generate_brick method
-        # Local BrickGPT uses max_new_tokens=10 for single brick generation
+        # Generation parameters - use BrickGPT defaults
         generation_params = {
-            "max_new_tokens": 10,  # EXACTLY like local BrickGPT
+            "max_new_tokens": parameters.get("max_new_tokens", 10),
             "temperature": parameters.get("temperature", 0.6),
             "top_k": parameters.get("top_k", 20),
             "top_p": parameters.get("top_p", 1.0),
             "pad_token_id": self.tokenizer.pad_token_id,
-            "do_sample": True,  # EXACTLY like local LLM
-            "num_return_sequences": 1,  # EXACTLY like local LLM
+            "do_sample": True,
+            "num_return_sequences": 1,
             "return_dict_in_generate": True,
         }
         
-        # Generate exactly like the local LLM class
+        # Generate
         with torch.no_grad():
             output_dict = self.model.generate(
                 input_ids,
@@ -160,17 +85,12 @@ class EndpointHandler:
                 **generation_params
             )
         
-        # Decode EXACTLY like local BrickGPT does
+        # Extract new tokens and decode EXACTLY like local LLM
         input_length = input_ids.shape[1]
         result_ids = output_dict['sequences'][0][input_length:]
         
-        # Local BrickGPT uses skip_special_tokens=True in generate_brick methods
-        generated_text = self.tokenizer.decode(result_ids, skip_special_tokens=True)
-        
-        # Extract LEGO instructions (same as before)
-        lego_instructions = self.extract_lego_instructions(generated_text)
+        # CRITICAL: Decode exactly like local LLM (no skip_special_tokens parameter)
+        generated_text = self.tokenizer.decode(result_ids)
         
-        return [{
-            "generated_text": generated_text,
-            "lego_instructions": lego_instructions
-        }]
+        # Return in format RemoteLLM expects
+        return [{"generated_text": generated_text}]