Update handler.py

handler.py

@@ -100,30 +100,73 @@ class EndpointHandler:
         Returns:
             PIL Image object or None if something goes wrong
         """
+        if not image_input or not isinstance(image_input, str):
+            print("❌ Invalid image input provided")
+            return None
+            
         try:
             # Check if it's a URL (starts with http/https)
-            if isinstance(image_input, str) and (image_input.startswith('http://') or image_input.startswith('https://')):
+            if image_input.startswith(('http://', 'https://')):
                 print(f"🌐 Fetching image from URL: {image_input[:50]}...")
-                response = requests.get(image_input, timeout=10)
+                headers = {
+                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
+                }
+                response = requests.get(image_input, timeout=15, headers=headers)
                 response.raise_for_status()
+                
+                # Verify it's actually an image
+                if not response.headers.get('content-type', '').startswith('image/'):
+                    print(f"⚠️ URL doesn't seem to point to an image: {response.headers.get('content-type')}")
+                
                 image = Image.open(BytesIO(response.content)).convert('RGB')
-                print("✅ Image downloaded successfully!")
+                print(f"✅ Image downloaded successfully! Size: {image.size}")
                 return image
             
-            # Must be base64 then
-            elif isinstance(image_input, str):
-                print("🔍 Decoding base64 image...")
-                # Remove the data URL prefix if it exists
-                if "base64," in image_input:
-                    image_input = image_input.split("base64,")[1]
-                
-                image_data = base64.b64decode(image_input)
+            # Handle base64 images
+            else:
+                print("🔍 Processing base64 image...")
+                base64_data = image_input
+                
+                # Remove data URL prefix if it exists (data:image/jpeg;base64,...)
+                if image_input.startswith('data:'):
+                    if 'base64,' in image_input:
+                        base64_data = image_input.split('base64,')[1]
+                    else:
+                        print("❌ Invalid data URL format - missing base64 encoding")
+                        return None
+                
+                # Clean up any whitespace
+                base64_data = base64_data.strip().replace('\n', '').replace('\r', '').replace(' ', '')
+                
+                # Validate base64 format
+                try:
+                    # Add padding if necessary
+                    missing_padding = len(base64_data) % 4
+                    if missing_padding:
+                        base64_data += '=' * (4 - missing_padding)
+                    
+                    image_data = base64.b64decode(base64_data, validate=True)
+                except Exception as decode_error:
+                    print(f"❌ Invalid base64 encoding: {decode_error}")
+                    return None
+                
+                # Verify it's a valid image
+                if len(image_data) < 100:  # Too small to be a real image
+                    print("❌ Decoded data too small to be a valid image")
+                    return None
+                
                 image = Image.open(BytesIO(image_data)).convert('RGB')
-                print("✅ Image decoded successfully!")
+                print(f"✅ Base64 image decoded successfully! Size: {image.size}")
                 return image
                 
+        except requests.exceptions.Timeout:
+            print("❌ Request timeout - image URL took too long to respond")
+            return None
+        except requests.exceptions.RequestException as e:
+            print(f"❌ Network error while fetching image: {e}")
+            return None
         except Exception as e:
-            print(f"❌ Couldn't process the image: {e}")
+            print(f"❌ Error processing image: {e}")
             return None
         
         return None
@@ -154,15 +197,20 @@ class EndpointHandler:
             
             if isinstance(inputs, dict):
                 # Dictionary input - check for text and image
-                text = inputs.get("text", inputs.get("prompt", str(inputs)))
+                # Support multiple text field names: query, text, prompt
+                text = inputs.get("query", inputs.get("text", inputs.get("prompt", "")))
                 
                 # Check for image in various formats
                 image_input = inputs.get("image", inputs.get("image_url", inputs.get("image_base64", None)))
                 if image_input:
                     image = self.process_image_input(image_input)
                     if image:
-                        # For now, we'll add a note about the image since we're text-only
-                        text = f"[Image provided - {image.size[0]}x{image.size[1]} pixels] {text}"
+                        print(f"✅ Image processed successfully: {image.size[0]}x{image.size[1]} pixels")
+                        # Add image context to the prompt for better processing
+                        if text:
+                            text = f"<image>\nUser query: {text}"
+                        else:
+                            text = "<image>\nAnalyze this medical image."
             else:
                 # Simple string input
                 text = str(inputs)
@@ -172,29 +220,52 @@ class EndpointHandler:
             
             # Get generation parameters with sensible defaults
             parameters = data.get("parameters", {})
-            max_new_tokens = min(parameters.get("max_new_tokens", 256), 1024)
-            temperature = parameters.get("temperature", 0.7)
-            top_p = parameters.get("top_p", 0.95)
-            do_sample = parameters.get("do_sample", True)
-            repetition_penalty = parameters.get("repetition_penalty", 1.0)
+            max_new_tokens = min(parameters.get("max_new_tokens", 512), 2048)  # Increased default
+            temperature = max(0.01, min(parameters.get("temperature", 0.2), 2.0))  # Clamp temperature
+            top_p = max(0.01, min(parameters.get("top_p", 0.9), 1.0))  # Clamp top_p
+            do_sample = parameters.get("do_sample", temperature > 0.01)  # Auto-set based on temperature
+            repetition_penalty = max(1.0, min(parameters.get("repetition_penalty", 1.05), 2.0))  # Clamp penalty
+            stop_sequences = parameters.get("stop", ["</s>"])  # Support stop sequences
+            return_full_text = parameters.get("return_full_text", False)
+            
+            print(f"🎛️ Generation params: max_tokens={max_new_tokens}, temp={temperature}, top_p={top_p}, rep_penalty={repetition_penalty}")
             
             # Using pipeline? Let's go!
             if self.use_pipeline:
-                result = self.pipe(
-                    text,
-                    max_new_tokens=max_new_tokens,
-                    temperature=temperature,
-                    top_p=top_p,
-                    do_sample=do_sample,
-                    repetition_penalty=repetition_penalty,
-                    return_full_text=False  # Just the new stuff, not the input
-                )
+                generation_kwargs = {
+                    "max_new_tokens": max_new_tokens,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "do_sample": do_sample,
+                    "repetition_penalty": repetition_penalty,
+                    "return_full_text": return_full_text
+                }
                 
-                # Pipeline returns a list, let's handle it
+                # Add stop sequences if supported
+                if stop_sequences and stop_sequences != ["</s>"]:
+                    generation_kwargs["stop_sequence"] = stop_sequences[0]  # Most pipelines support single stop
+                
+                result = self.pipe(text, **generation_kwargs)
+                
+                # Pipeline returns a list, let's handle it properly
                 if isinstance(result, list) and len(result) > 0:
-                    return [{"generated_text": result[0].get("generated_text", "")}]
+                    generated_text = result[0].get("generated_text", "")
+                    # Clean up any stop sequences that might remain
+                    for stop_seq in stop_sequences:
+                        if generated_text.endswith(stop_seq):
+                            generated_text = generated_text[:-len(stop_seq)].rstrip()
+                    
+                    return [{
+                        "generated_text": generated_text,
+                        "model": "PULSE-7B",
+                        "processing_method": "pipeline"
+                    }]
                 else:
-                    return [{"generated_text": str(result)}]
+                    return [{
+                        "generated_text": str(result),
+                        "model": "PULSE-7B", 
+                        "processing_method": "pipeline"
+                    }]
             
             # Manual generation mode
             else:
@@ -203,7 +274,7 @@ class EndpointHandler:
                     text,
                     return_tensors="pt",
                     truncation=True,
-                    max_length=2048
+                    max_length=4096  # Increased context length
                 )
                 
                 input_ids = encoded["input_ids"].to(self.device)
@@ -211,19 +282,33 @@ class EndpointHandler:
                 if attention_mask is not None:
                     attention_mask = attention_mask.to(self.device)
                 
+                # Prepare stop token IDs
+                stop_token_ids = []
+                if stop_sequences:
+                    for stop_seq in stop_sequences:
+                        stop_tokens = self.tokenizer.encode(stop_seq, add_special_tokens=False)
+                        if stop_tokens:
+                            stop_token_ids.extend(stop_tokens)
+                
                 # Generate the response
                 with torch.no_grad():
-                    outputs = self.model.generate(
-                        input_ids,
-                        attention_mask=attention_mask,
-                        max_new_tokens=max_new_tokens,
-                        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
-                    )
+                    generation_kwargs = {
+                        "input_ids": input_ids,
+                        "attention_mask": attention_mask,
+                        "max_new_tokens": max_new_tokens,
+                        "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
+                    }
+                    
+                    # Add stop token IDs if we have them
+                    if stop_token_ids:
+                        generation_kwargs["eos_token_id"] = stop_token_ids + [self.tokenizer.eos_token_id]
+                    
+                    outputs = self.model.generate(**generation_kwargs)
                 
                 # Decode only the new tokens (not the input)
                 generated_ids = outputs[0][input_ids.shape[-1]:]
@@ -233,13 +318,24 @@ class EndpointHandler:
                     clean_up_tokenization_spaces=True
                 )
                 
-                return [{"generated_text": generated_text}]
+                # Clean up any remaining stop sequences
+                for stop_seq in stop_sequences:
+                    if generated_text.endswith(stop_seq):
+                        generated_text = generated_text[:-len(stop_seq)].rstrip()
+                
+                return [{
+                    "generated_text": generated_text.strip(),
+                    "model": "PULSE-7B",
+                    "processing_method": "manual"
+                }]
             
         except Exception as e:
-            error_msg = f"Something went wrong during generation: {str(e)}"
+            error_msg = f"Generation error: {str(e)}"
             print(f"❌ {error_msg}")
             return [{
                 "generated_text": "",
                 "error": error_msg,
-                "handler": "Ubden® Team Enhanced Handler"
+                "model": "PULSE-7B",
+                "handler": "Ubden® Team Enhanced Handler",
+                "success": False
             }]