Update llama_customized_models.py

llama_customized_models.py

@@ -129,17 +129,22 @@ class LlamaForCausalLMWithNumericalEmbedding(LlamaForCausalLM):
         b, l = input_ids.size()
         assert len(properties) == b, "The number of properties should be equal to the batch size."
         assert len(properties_index) == b, "The number of properties_index should be equal to the batch size."
-        
+
+        print(input_ids, "input_ids")
         embeddings = self.model.embed_tokens(input_ids)
+        print(embeddings, "embeddings")
         
         for i, (props, props_index, embeds) in enumerate(zip(properties, properties_index, embeddings)):
             assert len(props) == len(props_index), "The number of properties should be equal to the number of properties_index."
             props = torch.tensor(props, device=embeds.device, dtype=torch.float32).unsqueeze(1)
+            print(props, "props")
             num_embeds = self.numerical_embedding(props)
+            print(num_embeds, "num_embeds")
             if len(props_index) > 0:
                 assert embeddings[i, props_index, :].shape == num_embeds.shape, "The shape of the embeddings and the numerical embeddings should be the same."
                 embeddings[i, props_index, :] = num_embeds
-
+            print(embeddings, "embedding_after")
+            
         return super().forward(
             input_ids=None,
             attention_mask=attention_mask,