Update mappings

Author: BenConstable9

deploy_ai_search_indexes/src/deploy_ai_search_indexes/ai_search.py

@@ -196,7 +196,7 @@ def get_data_source(self) -> SearchIndexerDataSourceConnection:
 
         return data_source_connection
 
-    def get_mark_up_cleaner_skill(self, context, source) -> WebApiSkill:
+    def get_mark_up_cleaner_skill(self, chunk_by_page: False) -> WebApiSkill:
         """Get the custom skill for data cleanup.
 
         Args:
@@ -215,29 +215,47 @@ def get_mark_up_cleaner_skill(self, context, source) -> WebApiSkill:
             batch_size = 16
             degree_of_parallelism = 16
 
-        mark_up_cleaner_skill_inputs = [
-            InputFieldMappingEntry(name="chunk", source=source),
-            InputFieldMappingEntry(
-                name="figure_storage_prefix", source="/document/metadata_storage_path"
-            ),
-        ]
+        if chunk_by_page:
+            mark_up_cleaner_context = "/document/page_wise_layout/*"
+            inputs = [
+                InputFieldMappingEntry(
+                    name="chunk", source="/document/page_wise_layout/*/merged_content"
+                ),
+                InputFieldMappingEntry(
+                    name="figures",
+                    source="/document/page_wise_layout/*/figures/*/updated_figure",
+                ),
+            ]
+        else:
+            mark_up_cleaner_context = "/document/chunk_mark_ups/*"
+            inputs = [
+                InputFieldMappingEntry(
+                    name="chunk", source="/document/chunk_mark_ups/*"
+                ),
+                InputFieldMappingEntry(
+                    name="figures", source="/document/layout/figures/*/updated_figure"
+                ),
+            ]
 
         mark_up_cleaner_skill_outputs = [
-            OutputFieldMappingEntry(name="cleaned_chunk", target_name="cleaned_chunk"),
-            OutputFieldMappingEntry(name="chunk", target_name="chunk"),
-            OutputFieldMappingEntry(name="sections", target_name="sections"),
+            OutputFieldMappingEntry(name="chunk_cleaned", target_name="chunk_cleaned"),
+            OutputFieldMappingEntry(
+                name="chunk_sections", target_name="chunk_sections"
+            ),
+            OutputFieldMappingEntry(name="chunk_mark_up", target_name="chunk_mark_up"),
+            OutputFieldMappingEntry(name="chunk_figures", target_name="chunk_figures"),
         ]
 
         mark_up_cleaner_skill = WebApiSkill(
             name="Mark Up Cleaner Skill",
             description="Skill to clean the data before sending to embedding",
-            context=context,
+            context=mark_up_cleaner_context,
             uri=self.environment.get_custom_skill_function_url("mark_up_cleaner"),
             timeout="PT230S",
             batch_size=batch_size,
             degree_of_parallelism=degree_of_parallelism,
             http_method="POST",
-            inputs=mark_up_cleaner_skill_inputs,
+            inputs=inputs,
             outputs=mark_up_cleaner_skill_outputs,
         )
 
@@ -255,8 +273,6 @@ def get_mark_up_cleaner_skill(self, context, source) -> WebApiSkill:
 
     def get_semantic_chunker_skill(
         self,
-        context,
-        source,
         num_surrounding_sentences: int = 1,
         similarity_threshold: float = 0.8,
         max_chunk_tokens: int = 200,
@@ -284,17 +300,17 @@ def get_semantic_chunker_skill(
             degree_of_parallelism = 16
 
         semantic_text_chunker_skill_inputs = [
-            InputFieldMappingEntry(name="content", source=source)
+            InputFieldMappingEntry(name="content", source="/document/merged_content")
         ]
 
         semantic_text_chunker_skill_outputs = [
-            OutputFieldMappingEntry(name="chunks", target_name="chunks"),
+            OutputFieldMappingEntry(name="chunks", target_name="chunk_mark_ups"),
         ]
 
         semantic_text_chunker_skill = WebApiSkill(
             name="Semantic Chunker Skill",
             description="Skill to clean the data before sending to embedding",
-            context=context,
+            context="/document",
             uri=self.environment.get_custom_skill_function_url("semantic_text_chunker"),
             timeout="PT230S",
             batch_size=batch_size,
@@ -345,7 +361,9 @@ def get_layout_analysis_skill(
 
         if chunk_by_page:
             output = [
-                OutputFieldMappingEntry(name="layout", target_name="page_wise_layout")
+                OutputFieldMappingEntry(
+                    name="page_wise_layout", target_name="page_wise_layout"
+                )
             ]
         else:
             output = [OutputFieldMappingEntry(name="layout", target_name="layout")]
@@ -383,7 +401,7 @@ def get_layout_analysis_skill(
 
         return layout_analysis_skill
 
-    def get_figure_analysis_skill(self, figure_source) -> WebApiSkill:
+    def get_figure_analysis_skill(self, chunk_by_page=False) -> WebApiSkill:
         """Get the custom skill for figure analysis.
 
         Args:
@@ -406,16 +424,32 @@ def get_figure_analysis_skill(self, figure_source) -> WebApiSkill:
             OutputFieldMappingEntry(name="updated_figure", target_name="updated_figure")
         ]
 
+        if chunk_by_page:
+            figure_context = "/document/page_wise_layout/*"
+            inputs = [
+                InputFieldMappingEntry(
+                    name="figure", source="/document/page_wise_layout/*/figures/*"
+                )
+            ]
+        else:
+            figure_context = "/document/layout/figures/*"
+
+            inputs = [
+                InputFieldMappingEntry(
+                    name="figure", source="/document/layout/figures/*"
+                )
+            ]
+
         figure_analysis_skill = WebApiSkill(
             name="Figure Analysis Skill",
             description="Skill to generate figure analysis",
-            context=figure_source,
+            context=figure_context,
             uri=self.environment.get_custom_skill_function_url("figure_analysis"),
             timeout="PT230S",
             batch_size=batch_size,
             degree_of_parallelism=degree_of_parallelism,
             http_method="POST",
-            inputs=[InputFieldMappingEntry(name="figure", source=figure_source)],
+            inputs=inputs,
             outputs=output,
         )
 
@@ -431,7 +465,7 @@ def get_figure_analysis_skill(self, figure_source) -> WebApiSkill:
 
         return figure_analysis_skill
 
-    def get_layout_and_figure_merger_skill(self, figure_source) -> WebApiSkill:
+    def get_layout_and_figure_merger_skill(self, chunk_by_page=False) -> WebApiSkill:
         """Get the custom skill for layout and figure merger.
 
         Args:
@@ -450,22 +484,40 @@ def get_layout_and_figure_merger_skill(self, figure_source) -> WebApiSkill:
             batch_size = 1
             degree_of_parallelism = 8
 
-        output = [
-            OutputFieldMappingEntry(name="updated_figure", target_name="updated_figure")
-        ]
+        output = [OutputFieldMappingEntry(name="content", target_name="merged_content")]
+        if chunk_by_page:
+            merger_context = "/document/page_wise_layout/*"
+            inputs = [
+                InputFieldMappingEntry(
+                    name="layout", source="/document/page_wise_layout/*"
+                ),
+                InputFieldMappingEntry(
+                    name="figures",
+                    source="/document/page_wise_layout/*/figures/*/updated_figure",
+                ),
+            ]
+        else:
+            merger_context = "/document/layout"
+
+            inputs = [
+                InputFieldMappingEntry(name="layout", source="/document/layout"),
+                InputFieldMappingEntry(
+                    name="figures", source="/document/layout/figures/*/updated_figure"
+                ),
+            ]
 
         figure_analysis_skill = WebApiSkill(
             name="Layout and Figure Merger Skill",
             description="Skill to merge layout and figure analysis",
-            context=figure_source,
+            context=merger_context,
             uri=self.environment.get_custom_skill_function_url(
                 "layout_and_figure_merger"
             ),
             timeout="PT230S",
             batch_size=batch_size,
             degree_of_parallelism=degree_of_parallelism,
             http_method="POST",
-            inputs=[InputFieldMappingEntry(name="figure", source=figure_source)],
+            inputs=inputs,
             outputs=output,
         )
 

deploy_ai_search_indexes/src/deploy_ai_search_indexes/image_processing.py

@@ -181,21 +181,27 @@ def get_skills(self) -> list:
 
         layout_skill = self.get_layout_analysis_skill(self.enable_page_by_chunking)
 
-        figure_skill = self.get_figure_analysis_skill()
+        figure_skill = self.get_figure_analysis_skill(self.enable_page_by_chunking)
 
-        merger_skill = self.get_layout_and_figure_merger_skill()
-
-        text_split_skill = self.get_semantic_chunker_skill(
-            "/document", "/document/extracted_content/content"
+        merger_skill = self.get_layout_and_figure_merger_skill(
+            self.enable_page_by_chunking
         )
 
+        text_split_skill = self.get_semantic_chunker_skill(self.enable_page_by_chunking)
+
         mark_up_cleaner_skill = self.get_mark_up_cleaner_skill(
-            "/document/chunks/*", "/document/chunks/*/content"
+            self.enable_page_by_chunking
         )
 
-        embedding_skill = self.get_vector_skill(
-            "/document/chunks/*", "/document/chunks/*/cleaned_chunk"
-        )
+        if self.enable_page_by_chunking:
+            embedding_skill = self.get_vector_skill(
+                "/document/page_wise_layout/*",
+                "/document/page_wise_layout/*/chunk_cleaned",
+            )
+        else:
+            embedding_skill = self.get_vector_skill(
+                "/document/chunk_mark_ups/*", "/document/chunk_mark_ups/*/chunk_cleaned"
+            )
 
         if self.enable_page_by_chunking:
             skills = [
@@ -219,41 +225,64 @@ def get_skills(self) -> list:
 
     def get_index_projections(self) -> SearchIndexerIndexProjection:
         """This function returns the index projections for rag document."""
-        mappings = [
-            InputFieldMappingEntry(name="Chunk", source="/document/chunks/*/chunk"),
-            InputFieldMappingEntry(
-                name="ChunkEmbedding",
-                source="/document/chunks/*/vector",
-            ),
-            InputFieldMappingEntry(name="Title", source="/document/Title"),
-            InputFieldMappingEntry(name="SourceUri", source="/document/SourceUri"),
-            InputFieldMappingEntry(
-                name="Sections", source="/document/chunks/*/sections"
-            ),
-            InputFieldMappingEntry(
-                name="Figures",
-                source_context="/document/chunks/*/chunk_figures/*",
-            ),
-            InputFieldMappingEntry(
-                name="DateLastModified", source="/document/DateLastModified"
-            ),
-        ]
 
         if self.enable_page_by_chunking:
-            mappings.extend(
-                [
-                    InputFieldMappingEntry(
-                        name="PageNumber", source="/document/chunks/*/pageNumber"
-                    )
-                ]
-            )
+            source_context = "/document/page_wise_layout/*"
+            mappings = [
+                InputFieldMappingEntry(
+                    name="Chunk", source="/document/page_wise_layout/*/chunk_mark_up"
+                ),
+                InputFieldMappingEntry(
+                    name="ChunkEmbedding",
+                    source="/document/page_wise_layout/*/vector",
+                ),
+                InputFieldMappingEntry(name="Title", source="/document/Title"),
+                InputFieldMappingEntry(name="SourceUri", source="/document/SourceUri"),
+                InputFieldMappingEntry(
+                    name="Sections",
+                    source="/document/page_wise_layout/*/chunk_sections",
+                ),
+                InputFieldMappingEntry(
+                    name="Figures",
+                    source_context="/document/page_wise_layout/*/chunk_figures/*",
+                ),
+                InputFieldMappingEntry(
+                    name="DateLastModified", source="/document/DateLastModified"
+                ),
+                InputFieldMappingEntry(
+                    name="PageNumber", source="/document/page_wise_layout/*/#"
+                ),
+            ]
+        else:
+            source_context = "/document/chunks/*"
+            mappings = [
+                InputFieldMappingEntry(
+                    name="Chunk", source="/document/chunk_mark_ups/*/chunk_mark_up"
+                ),
+                InputFieldMappingEntry(
+                    name="ChunkEmbedding",
+                    source="/document/chunk_mark_ups/*/vector",
+                ),
+                InputFieldMappingEntry(name="Title", source="/document/Title"),
+                InputFieldMappingEntry(name="SourceUri", source="/document/SourceUri"),
+                InputFieldMappingEntry(
+                    name="Sections", source="/document/chunk_mark_ups/*/chunk_sections"
+                ),
+                InputFieldMappingEntry(
+                    name="Figures",
+                    source_context="/document/chunk_mark_ups/*/chunk_figures/*",
+                ),
+                InputFieldMappingEntry(
+                    name="DateLastModified", source="/document/DateLastModified"
+                ),
+            ]
 
         index_projections = SearchIndexerIndexProjection(
             selectors=[
                 SearchIndexerIndexProjectionSelector(
                     target_index_name=self.index_name,
                     parent_key_field_name="Id",
-                    source_context="/document/chunks/*",
+                    source_context=source_context,
                     mappings=mappings,
                 ),
             ],