fix added

SkinGPT.py

@@ -229,6 +229,10 @@ class SkinGPT4(nn.Module):
 
         print(f"Aligned features : {image_embeds}")
         print(f"\n Images embeddings shape : {image_embeds.shape} \n Llama config hidden size : {self.llama.config.hidden_size}")
+
+        print(
+            f"\n[VALIDATION] Visual embeds - Mean: {image_embeds.mean().item():.4f}, Std: {image_embeds.std().item():.4f}")
+
         if image_embeds.shape[-1] != self.llama.config.hidden_size:
             raise ValueError(
                 f"Feature dimension mismatch. "
@@ -238,24 +242,20 @@ class SkinGPT4(nn.Module):
 
 
         # prompt = (
-        #     "### Instruction: <Img><ImageHere></Img> "
+        #     "### Instruction: <Img><IMAGE></Img> "
         #     "Could you describe the skin condition in this image? "
         #     "Please provide a detailed analysis including possible diagnoses. "
         #     "### Response:"
         # )
 
-        prompt = """### Skin Diagnosis Protocol ###
+        prompt = """### Skin Diagnosis Analysis ###
         <IMAGE>
-        Patient Presentation: [Describe visible symptoms]
-        Primary Differential Diagnosis:
-        1. 
-        2. 
-        3. 
-        Recommended Diagnostic Tests:
-        - 
-        Treatment Options:
-        - 
-        <|endoftext|>"""
+        Describe the skin condition shown and provide:
+        1. Primary diagnosis (with confidence)
+        2. Three differential diagnoses
+        3. Recommended tests
+        4. Treatment options"""
+
 
         print(f"\n[DEBUG] Raw Prompt:\n{prompt}")
 
@@ -265,7 +265,13 @@ class SkinGPT4(nn.Module):
             padding_side="right"
         )
         # self.tokenizer.add_special_tokens({'additional_special_tokens': ['<Img>', '</Img>', '<ImageHere>']})
-        self.tokenizer.add_tokens(["<IMAGE>"])
+        num_added = self.tokenizer.add_special_tokens({
+            'additional_special_tokens': ['<IMAGE>']
+        })
+
+        if num_added == 0:
+            raise ValueError("Failed to add <IMAGE> token!")
+
         self.llama.resize_token_embeddings(len(self.tokenizer))
 
         inputs = self.tokenizer(prompt, return_tensors="pt").to(images.device)
@@ -273,37 +279,37 @@ class SkinGPT4(nn.Module):
         print(f"\n[DEBUG] Tokenized input IDs:\n{inputs.input_ids}")
         print(f"[DEBUG] Special token positions: {self.tokenizer.all_special_tokens}")
 
-        # image_token_id = self.tokenizer.convert_tokens_to_ids("<ImageHere>")
-        image_token_id = self.tokenizer.convert_tokens_to_ids("<IMAGE>")
-
-
         # Prepare embeddings
         input_embeddings = self.llama.model.embed_tokens(inputs.input_ids)
-        visual_embeds = image_embeds.mean(dim=1, keepdim=True)
+        visual_embeds = image_embeds.mean(dim=1)
+
+        # image_token_id = self.tokenizer.convert_tokens_to_ids("<ImageHere>")
+        image_token_id = self.tokenizer.convert_tokens_to_ids("<IMAGE>")
+        replace_positions = (inputs.input_ids == image_token_id).nonzero()
 
-        visual_embeds = F.layer_norm(visual_embeds, [visual_embeds.size(-1)])
-        image_token_pos = torch.where(inputs.input_ids == image_token_id)
+        if len(replace_positions) == 0:
+            raise ValueError("No <IMAGE> tokens found in prompt!")
 
-        if len(image_token_pos[0]) == 0:
+        if len(replace_positions[0]) == 0:
             raise ValueError("Image token not found in prompt")
 
-        print(f"\n[DEBUG] Image token found at position: {image_token_pos}")
+        print(f"\n[DEBUG] Image token found at position: {replace_positions}")
 
-        for pos in zip(*image_token_pos):
-            input_embeddings[pos] = visual_embeds[0, 0]
 
         print(f"\n[DEBUG] Before replacement:")
         print(f"Text embeddings shape: {input_embeddings.shape}")
         print(f"Visual embeddings shape: {visual_embeds.shape}")
-        print(f"Image token embedding (before):\n{input_embeddings[0, image_token_pos[1], :5]}...")
+        print(f"Image token embedding (before):\n{input_embeddings[0, replace_positions[1], :5]}...")
 
+        for pos in replace_positions:
+            input_embeddings[0, pos[1]] = visual_embeds[0]
 
-        if visual_embeds.dtype != input_embeddings.dtype:
-            visual_embeds = visual_embeds.to(input_embeddings.dtype)
+        # if visual_embeds.dtype != input_embeddings.dtype:
+        #     visual_embeds = visual_embeds.to(input_embeddings.dtype)
         # input_embeddings[image_token_pos] = visual_embeds
 
         print(f"\n[DEBUG] After replacement:")
-        print(f"Image token embedding (after):\n{input_embeddings[0, image_token_pos[1], :5]}...")
+        print(f"Image token embedding (after):\n{input_embeddings[0, replace_positions[1], :5]}...")
 
         # outputs = self.llama.generate(
         #     inputs_embeds=input_embeddings,
@@ -340,7 +346,7 @@ class SkinGPT4(nn.Module):
         outputs = self.llama.generate(
             inputs_embeds=input_embeddings,
             max_new_tokens=max_new_tokens,
-            temperature=0.3,
+            temperature=0.7,
             top_k=40,
             top_p=0.9,
             repetition_penalty=1.1,