correction of columns of datasets

src/api/models/embedding_models.py

@@ -69,3 +69,4 @@ class SearchEmbeddingRequest(BaseModel):
 
 class ResetEmbeddingsRequest(BaseModel):
     dataset_name: str = "re-mind/product_type_embedding"
+    target_column: str = "product_type"

src/api/services/postgresql_service.py

@@ -13,26 +13,49 @@ class PostgresqlService:
         query_mapping = {
             "re-mind/product_type_embedding": {
                 "column": "type",
+                "target_column": "product_type",
                 "table": "product_producttype",
+                "optional": {
+                    "column": "name",
+                    "target_column": "product_category",
+                    "table": "product_category",
+                    "foreign_key": "category_id",
+                }
             },
             "re-mind/marketplace_name_embedding": {
                 "column": "name",
+                "target_column": "marketplace_name",
                 "table": "invoice_marketplace",
             },
             "re-mind/manufacturer_name_embedding": {
                 "column": "name",
+                "target_column": "manufacturer_name",
                 "table": "product_manufacturer",
             },
             "re-mind/seller_name_embedding": {
                 "column": "name",
+                "target_column": "seller_name",
                 "table": "product_seller",
             }
         }
 
         db_table = query_mapping[dataset]["table"]
         column = query_mapping[dataset]["column"]
-        query = f"SELECT {column} FROM {db_table}"
+        target_column = query_mapping[dataset]["target_column"]
+        optional = query_mapping[dataset].get("optional")
+        optional_query_1 = ""
+        optional_query_2 = ""
+        if optional:
+            optional_table = optional["table"]
+            optional_column = optional["column"]
+            optional_target_column = optional["target_column"]
+            optional_foreign_key = optional["foreign_key"]
 
+            optional_query_1 = f", {optional_table}.{optional_column} AS {optional_target_column}"
+            optional_query_2 = f"INNER JOIN {optional_table} ON {db_table}.{optional_foreign_key} = {optional_table}.id"
+        query = f"SELECT {db_table}.{column} AS {target_column} {optional_query_1} FROM {db_table} {optional_query_2}"
+
+        logger.info(query)
         results = await self.db.fetch(query)
 
         return results

src/main.py

@@ -327,8 +327,7 @@ async def reset_embeddings(
 
         # Generation of embeddings for each row
         dataset = Dataset.from_dict(results)
-        target_column = "type" if request.dataset_name == "re-mind/product_type_embedding" else "name"
-        dataset_embedded = await embedding_service.create_embeddings(dataset, target_column, "embedding")
+        dataset_embedded = await embedding_service.create_embeddings(dataset, request.target_column, "embedding")
         # Embeddings up-to-date with database will overwrite old dataset
         await huggingface_service.push_to_hub(dataset_embedded, request.dataset_name)