add fewshots

handler.py

@@ -2,6 +2,8 @@ from typing import Dict, List, Any
 import torch
 import re
 import os
+import json
+from pathlib import Path
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
@@ -25,6 +27,56 @@ 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"""
@@ -46,89 +98,102 @@ class EndpointHandler:
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
-        Process inference request
-        data args:
-            inputs (:obj:`str` or :obj:`Dict`): The input text or messages
-            parameters (:obj:`Dict`, optional): Parameters for generation
+        Process inference request exactly like BrickGPT does
         """
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
         
         # Handle different input formats that BrickGPT sends
         if isinstance(inputs, dict) and "messages" in inputs:
-            # BrickGPT format: {"messages": [{"role": "system", ...}, {"role": "user", ...}]}
             messages = inputs["messages"]
         elif isinstance(inputs, list):
-            # Direct messages array: [{"role": "system", ...}, {"role": "user", ...}]
             messages = inputs
         elif isinstance(inputs, str):
-            # Plain string input - create default messages
             messages = [
                 {"role": "system", "content": "You are a helpful assistant."},
                 {"role": "user", "content": inputs}
             ]
         else:
-            # Fallback
             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 input using chat template
-        if hasattr(self.tokenizer, 'chat_template') and self.tokenizer.chat_template:            
-            if has_assistant:
-                # For continuation, use continue_final_message=True
-                formatted_input = self.tokenizer.apply_chat_template(
-                    messages, 
-                    tokenize=False, 
-                    continue_final_message=True
-                )
-            else:
-                # For new generation, add generation prompt
-                formatted_input = self.tokenizer.apply_chat_template(
-                    messages, 
-                    tokenize=False, 
-                    add_generation_prompt=True
-                )
+        # Extract the actual user instruction
+        user_content = ""
+        for msg in messages:
+            if msg.get("role") == "user":
+                content = msg["content"]
+                if "### Input:" in content:
+                    user_content = content.split("### Input:")[-1].strip()
+                else:
+                    user_content = content
+                break
+        
+        # Create the proper instruction format (use few_shot for better results)
+        if not has_assistant:
+            instruction = self.create_instruction_few_shot(user_content)
+            messages = [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": instruction}
+            ]
+        
+        # Format input using chat template exactly like BrickGPT
+        if has_assistant:
+            # For continuation, use continue_final_message=True
+            formatted_input = self.tokenizer.apply_chat_template(
+                messages, 
+                tokenize=False,
+                continue_final_message=True
+            )
         else:
-            # Fallback to direct input (last user message)
-            user_messages = [msg["content"] for msg in messages if msg.get("role") == "user"]
-            formatted_input = user_messages[-1] if user_messages else str(inputs)
+            # For new generation, add generation prompt
+            formatted_input = self.tokenizer.apply_chat_template(
+                messages, 
+                tokenize=False,
+                add_generation_prompt=True
+            )
         
-        # Default parameters optimized for BrickGPT
-        # BrickGPT generates SHORT responses (one brick at a time), not long descriptions
-        default_max_tokens = 20 if has_assistant else 200  # Much shorter for continuation
+        # Generation parameters that match BrickGPT's approach
+        default_max_tokens = 15 if has_assistant else 50  # BrickGPT generates one brick at a time
         generation_params = {
             "max_new_tokens": parameters.get("max_new_tokens", default_max_tokens),
-            "temperature": parameters.get("temperature", 0.6),
-            "top_k": parameters.get("top_k", 20),
-            "top_p": parameters.get("top_p", 1.0),
+            "temperature": parameters.get("temperature", 0.6),  # BrickGPT default
+            "top_k": parameters.get("top_k", 20),  # BrickGPT default
+            "top_p": parameters.get("top_p", 1.0),   # BrickGPT default
             "do_sample": parameters.get("do_sample", True),
             "pad_token_id": self.tokenizer.pad_token_id,
+            "return_dict_in_generate": True
         }
         
-        # Add stop tokens if provided
-        stop_tokens = parameters.get("stop", [])
-        if stop_tokens:
-            generation_params["stop_strings"] = stop_tokens
-        
         # Tokenize input
         input_ids = self.tokenizer(formatted_input, return_tensors="pt").input_ids.to(self.model.device)
+        attention_mask = torch.ones_like(input_ids)
         
-        # Generate
+        # Generate exactly like the local LLM class
         with torch.no_grad():
-            outputs = self.model.generate(input_ids, **generation_params)
+            output_dict = self.model.generate(
+                input_ids,
+                attention_mask=attention_mask,
+                **generation_params
+            )
         
-        # Decode output
-        generated_text = self.tokenizer.decode(
-            outputs[0][input_ids.shape[1]:], 
-            skip_special_tokens=True
-        )
+        # Decode exactly like the local LLM class  
+        input_length = input_ids.shape[1]
+        result_ids = output_dict['sequences'][0][input_length:]
+        generated_text = self.tokenizer.decode(result_ids, skip_special_tokens=True)
+        
+        # Clean up the generated text
+        generated_text = generated_text.strip()
+        
+        # Remove any trailing continuation artifacts
+        if generated_text.endswith("### Output:"):
+            generated_text = generated_text[:-11].strip()
         
         # Extract LEGO instructions
         lego_instructions = self.extract_lego_instructions(generated_text)
         
         return [{
             "generated_text": generated_text,
-            "lego_instructions": lego_instructions  # Fixed the field name
+            "lego_instructions": lego_instructions
         }]