upgini 1.1.315a3579.dev1__py3-none-any.whl → 1.1.316a1__py3-none-any.whl

upgini/http.py

@@ -39,17 +39,6 @@ from upgini.metadata import (
 from upgini.resource_bundle import bundle
 from upgini.utils.track_info import get_track_metrics
 
-# try:
-#     from importlib.metadata import version  # type: ignore
-
-#     __version__ = version("upgini")
-# except ImportError:
-#     try:
-#         from importlib_metadata import version  # type: ignore
-
-#         __version__ = version("upgini")
-#     except ImportError:
-#         __version__ = "Upgini wasn't installed"
 
 UPGINI_URL: str = "UPGINI_URL"
 UPGINI_API_KEY: str = "UPGINI_API_KEY"
@@ -452,18 +441,18 @@ class _RestClient:
                 content = file.read()
                 md5_hash.update(content)
                 digest = md5_hash.hexdigest()
-                metadata_with_md5 = metadata.copy(update={"checksumMD5": digest})
+                metadata_with_md5 = metadata.model_copy(update={"checksumMD5": digest})
 
             digest_sha256 = hashlib.sha256(
                 pd.util.hash_pandas_object(pd.read_parquet(file_path, engine="fastparquet")).values
             ).hexdigest()
-            metadata_with_md5 = metadata_with_md5.copy(update={"digest": digest_sha256})
+            metadata_with_md5 = metadata_with_md5.model_copy(update={"digest": digest_sha256})
 
             with open(file_path, "rb") as file:
                 files = {
                     "metadata": (
                         "metadata.json",
-                        metadata_with_md5.json(exclude_none=True).encode(),
+                        metadata_with_md5.model_dump_json(exclude_none=True).encode(),
                         "application/json",
                     ),
                     "tracking": (
@@ -471,13 +460,13 @@ class _RestClient:
                         dumps(track_metrics).encode(),
                         "application/json",
                     ),
-                    "metrics": ("metrics.json", metrics.json(exclude_none=True).encode(), "application/json"),
+                    "metrics": ("metrics.json", metrics.model_dump_json(exclude_none=True).encode(), "application/json"),
                     "file": (metadata_with_md5.name, file, "application/octet-stream"),
                 }
                 if search_customization is not None:
                     files["customization"] = (
                         "customization.json",
-                        search_customization.json(exclude_none=True).encode(),
+                        search_customization.model_dump_json(exclude_none=True).encode(),
                         "application/json",
                     )
                 additional_headers = {self.SEARCH_KEYS_HEADER_NAME: ",".join(self.search_keys_meaning_types(metadata))}
@@ -492,7 +481,7 @@ class _RestClient:
     def check_uploaded_file_v2(self, trace_id: str, file_upload_id: str, metadata: FileMetadata) -> bool:
         api_path = self.CHECK_UPLOADED_FILE_URL_FMT_V2.format(file_upload_id)
         response = self._with_unauth_retry(
-            lambda: self._send_post_req(api_path, trace_id, metadata.json(exclude_none=True))
+            lambda: self._send_post_req(api_path, trace_id, metadata.model_dump_json(exclude_none=True))
         )
         return bool(response)
 
@@ -506,11 +495,11 @@ class _RestClient:
     ) -> SearchTaskResponse:
         api_path = self.INITIAL_SEARCH_WITHOUT_UPLOAD_URI_FMT_V2.format(file_upload_id)
         files = {
-            "metadata": ("metadata.json", metadata.json(exclude_none=True).encode(), "application/json"),
-            "metrics": ("metrics.json", metrics.json(exclude_none=True).encode(), "application/json"),
+            "metadata": ("metadata.json", metadata.model_dump_json(exclude_none=True).encode(), "application/json"),
+            "metrics": ("metrics.json", metrics.model_dump_json(exclude_none=True).encode(), "application/json"),
         }
         if search_customization is not None:
-            files["customization"] = search_customization.json(exclude_none=True).encode()
+            files["customization"] = search_customization.model_dump_json(exclude_none=True).encode()
         additional_headers = {self.SEARCH_KEYS_HEADER_NAME: ",".join(self.search_keys_meaning_types(metadata))}
         response = self._with_unauth_retry(
             lambda: self._send_post_file_req_v2(
@@ -536,18 +525,18 @@ class _RestClient:
                 content = file.read()
                 md5_hash.update(content)
                 digest = md5_hash.hexdigest()
-                metadata_with_md5 = metadata.copy(update={"checksumMD5": digest})
+                metadata_with_md5 = metadata.model_copy(update={"checksumMD5": digest})
 
             digest_sha256 = hashlib.sha256(
                 pd.util.hash_pandas_object(pd.read_parquet(file_path, engine="fastparquet")).values
             ).hexdigest()
-            metadata_with_md5 = metadata_with_md5.copy(update={"digest": digest_sha256})
+            metadata_with_md5 = metadata_with_md5.model_copy(update={"digest": digest_sha256})
 
             with open(file_path, "rb") as file:
                 files = {
                     "metadata": (
                         "metadata.json",
-                        metadata_with_md5.json(exclude_none=True).encode(),
+                        metadata_with_md5.model_dump_json(exclude_none=True).encode(),
                         "application/json",
                     ),
                     "tracking": (
@@ -555,13 +544,13 @@ class _RestClient:
                         dumps(get_track_metrics(self.client_ip, self.client_visitorid)).encode(),
                         "application/json",
                     ),
-                    "metrics": ("metrics.json", metrics.json(exclude_none=True).encode(), "application/json"),
+                    "metrics": ("metrics.json", metrics.model_dump_json(exclude_none=True).encode(), "application/json"),
                     "file": (metadata_with_md5.name, file, "application/octet-stream"),
                 }
                 if search_customization is not None:
                     files["customization"] = (
                         "customization.json",
-                        search_customization.json(exclude_none=True).encode(),
+                        search_customization.model_dump_json(exclude_none=True).encode(),
                         "application/json",
                     )
 
@@ -585,11 +574,11 @@ class _RestClient:
     ) -> SearchTaskResponse:
         api_path = self.VALIDATION_SEARCH_WITHOUT_UPLOAD_URI_FMT_V2.format(file_upload_id, initial_search_task_id)
         files = {
-            "metadata": ("metadata.json", metadata.json(exclude_none=True).encode(), "application/json"),
-            "metrics": ("metrics.json", metrics.json(exclude_none=True).encode(), "application/json"),
+            "metadata": ("metadata.json", metadata.model_dump_json(exclude_none=True).encode(), "application/json"),
+            "metrics": ("metrics.json", metrics.model_dump_json(exclude_none=True).encode(), "application/json"),
         }
         if search_customization is not None:
-            files["customization"] = search_customization.json(exclude_none=True).encode()
+            files["customization"] = search_customization.model_dump_json(exclude_none=True).encode()
         additional_headers = {self.SEARCH_KEYS_HEADER_NAME: ",".join(self.search_keys_meaning_types(metadata))}
         response = self._with_unauth_retry(
             lambda: self._send_post_file_req_v2(
@@ -651,7 +640,7 @@ class _RestClient:
             with open(file_path, "rb") as file:
                 files = {
                     "file": (metadata.name, file, "application/octet-stream"),
-                    "metadata": ("metadata.json", metadata.json(exclude_none=True).encode(), "application/json"),
+                    "metadata": ("metadata.json", metadata.model_dump_json(exclude_none=True).encode(), "application/json"),
                 }
 
                 return self._send_post_file_req_v2(api_path, files)
@@ -661,12 +650,12 @@ class _RestClient:
     def get_search_file_metadata(self, search_task_id: str, trace_id: str) -> FileMetadata:
         api_path = self.SEARCH_FILE_METADATA_URI_FMT_V2.format(search_task_id)
         response = self._with_unauth_retry(lambda: self._send_get_req(api_path, trace_id))
-        return FileMetadata.parse_obj(response)
+        return FileMetadata.model_validate(response)
 
     def get_provider_search_metadata_v3(self, provider_search_task_id: str, trace_id: str) -> ProviderTaskMetadataV2:
         api_path = self.SEARCH_TASK_METADATA_FMT_V3.format(provider_search_task_id)
         response = self._with_unauth_retry(lambda: self._send_get_req(api_path, trace_id))
-        return ProviderTaskMetadataV2.parse_obj(response)
+        return ProviderTaskMetadataV2.model_validate(response)
 
     def get_current_transform_usage(self, trace_id) -> TransformUsage:
         track_metrics = get_track_metrics(self.client_ip, self.client_visitorid)

upgini/lazy_import.py

@@ -1,4 +1,6 @@
 import importlib
+import importlib.util
+import importlib.machinery
 
 
 class LazyImport:
@@ -10,7 +12,18 @@ class LazyImport:
 
     def _load(self):
         if self._module is None:
-            self._module = importlib.import_module(self.module_name)
+            # Load module and save link to it
+            spec = importlib.util.find_spec(self.module_name)
+            if spec is None:
+                raise ImportError(f"Module {self.module_name} not found")
+
+            # Create module
+            self._module = importlib.util.module_from_spec(spec)
+
+            # Execute module
+            spec.loader.exec_module(self._module)
+
+            # Get class from module
             self._class = getattr(self._module, self.class_name)
 
     def __call__(self, *args, **kwargs):

upgini/metadata.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Dict, List, Optional, Set
+from typing import Any, Dict, List, Optional, Set, Union
 
 from pydantic import BaseModel
 
@@ -113,6 +113,21 @@ class SearchKey(Enum):
         if meaning_type == FileColumnMeaningType.MSISDN_RANGE_TO:
             return SearchKey.MSISDN_RANGE_TO
 
+    @staticmethod
+    def find_key(search_keys: Dict[str, SearchKey], keys: Union[SearchKey, List[SearchKey]]) -> Optional[SearchKey]:
+        if isinstance(keys, SearchKey):
+            keys = [keys]
+        for col, key_type in search_keys.items():
+            if key_type in keys:
+                return col
+        return None
+
+    @staticmethod
+    def find_all_keys(search_keys: Dict[str, SearchKey], keys: Union[SearchKey, List[SearchKey]]) -> List[SearchKey]:
+        if isinstance(keys, SearchKey):
+            keys = [keys]
+        return [col for col, key_type in search_keys.items() if key_type in keys]
+
 
 class DataType(Enum):
     INT = "INT"
@@ -157,23 +172,23 @@ class FileMetricsInterval(BaseModel):
     date_cut: float
     count: float
     valid_count: float
-    avg_target: Optional[float]  # not for multiclass
-    avg_score_etalon: Optional[float]
+    avg_target: Optional[float] = None  # not for multiclass
+    avg_score_etalon: Optional[float] = None
 
 
 class FileMetrics(BaseModel):
     # etalon metadata
-    task_type: Optional[ModelTaskType]
-    label: Optional[ModelLabelType]
-    count: Optional[int]
-    valid_count: Optional[int]
-    valid_rate: Optional[float]
-    avg_target: Optional[float]
-    metrics_binary_etalon: Optional[BinaryTask]
-    metrics_regression_etalon: Optional[RegressionTask]
-    metrics_multiclass_etalon: Optional[MulticlassTask]
-    cuts: Optional[List[float]]
-    interval: Optional[List[FileMetricsInterval]]
+    task_type: Optional[ModelTaskType] = None
+    label: Optional[ModelLabelType] = None
+    count: Optional[int] = None
+    valid_count: Optional[int] = None
+    valid_rate: Optional[float] = None
+    avg_target: Optional[float] = None
+    metrics_binary_etalon: Optional[BinaryTask] = None
+    metrics_regression_etalon: Optional[RegressionTask] = None
+    metrics_multiclass_etalon: Optional[MulticlassTask] = None
+    cuts: Optional[List[float]] = None
+    interval: Optional[List[FileMetricsInterval]] = None
 
 
 class NumericInterval(BaseModel):
@@ -187,25 +202,25 @@ class FileColumnMetadata(BaseModel):
     dataType: DataType
     meaningType: FileColumnMeaningType
     minMaxValues: Optional[NumericInterval] = None
-    originalName: Optional[str]
+    originalName: Optional[str] = None
     # is this column contains keys from multiple key columns like msisdn1, msisdn2
     isUnnest: bool = False
     # list of original etalon key column names like msisdn1, msisdn2
-    unnestKeyNames: Optional[List[str]]
+    unnestKeyNames: Optional[List[str]] = None
 
 
 class FileMetadata(BaseModel):
     name: str
-    description: Optional[str]
+    description: Optional[str] = None
     columns: List[FileColumnMetadata]
     searchKeys: List[List[str]]
-    excludeFeaturesSources: Optional[List[str]]
-    hierarchicalGroupKeys: Optional[List[str]]
-    hierarchicalSubgroupKeys: Optional[List[str]]
-    taskType: Optional[ModelTaskType]
-    rowsCount: Optional[int]
-    checksumMD5: Optional[str]
-    digest: Optional[str]
+    excludeFeaturesSources: Optional[List[str]] = None
+    hierarchicalGroupKeys: Optional[List[str]] = None
+    hierarchicalSubgroupKeys: Optional[List[str]] = None
+    taskType: Optional[ModelTaskType] = None
+    rowsCount: Optional[int] = None
+    checksumMD5: Optional[str] = None
+    digest: Optional[str] = None
 
     def column_by_name(self, name: str) -> Optional[FileColumnMetadata]:
         for c in self.columns:
@@ -229,17 +244,17 @@ class FeaturesMetadataV2(BaseModel):
     source: str
     hit_rate: float
     shap_value: float
-    commercial_schema: Optional[str]
-    data_provider: Optional[str]
-    data_providers: Optional[List[str]]
-    data_provider_link: Optional[str]
-    data_provider_links: Optional[List[str]]
-    data_source: Optional[str]
-    data_sources: Optional[List[str]]
-    data_source_link: Optional[str]
-    data_source_links: Optional[List[str]]
-    doc_link: Optional[str]
-    update_frequency: Optional[str]
+    commercial_schema: Optional[str] = None
+    data_provider: Optional[str] = None
+    data_providers: Optional[List[str]] = None
+    data_provider_link: Optional[str] = None
+    data_provider_links: Optional[List[str]] = None
+    data_source: Optional[str] = None
+    data_sources: Optional[List[str]] = None
+    data_source_link: Optional[str] = None
+    data_source_links: Optional[List[str]] = None
+    doc_link: Optional[str] = None
+    update_frequency: Optional[str] = None
 
 
 class HitRateMetrics(BaseModel):
@@ -259,48 +274,48 @@ class ModelEvalSet(BaseModel):
 class BaseColumnMetadata(BaseModel):
     original_name: str
     hashed_name: str
-    ads_definition_id: Optional[str]
+    ads_definition_id: Optional[str] = None
     is_augmented: bool
 
 
 class GeneratedFeatureMetadata(BaseModel):
-    alias: Optional[str]
+    alias: Optional[str] = None
     formula: str
     display_index: str
     base_columns: List[BaseColumnMetadata]
-    operator_params: Optional[Dict[str, str]]
+    operator_params: Optional[Dict[str, str]] = None
 
 
 class ProviderTaskMetadataV2(BaseModel):
     features: List[FeaturesMetadataV2]
-    hit_rate_metrics: Optional[HitRateMetrics]
-    eval_set_metrics: Optional[List[ModelEvalSet]]
-    zero_hit_rate_search_keys: Optional[List[str]]
-    features_used_for_embeddings: Optional[List[str]]
-    shuffle_kfold: Optional[bool]
-    generated_features: Optional[List[GeneratedFeatureMetadata]]
+    hit_rate_metrics: Optional[HitRateMetrics] = None
+    eval_set_metrics: Optional[List[ModelEvalSet]] = None
+    zero_hit_rate_search_keys: Optional[List[str]] = None
+    features_used_for_embeddings: Optional[List[str]] = None
+    shuffle_kfold: Optional[bool] = None
+    generated_features: Optional[List[GeneratedFeatureMetadata]] = None
 
 
 class FeaturesFilter(BaseModel):
-    minImportance: Optional[float]
-    maxPSI: Optional[float]
-    maxCount: Optional[int]
-    selectedFeatures: Optional[List[str]]
+    minImportance: Optional[float] = None
+    maxPSI: Optional[float] = None
+    maxCount: Optional[int] = None
+    selectedFeatures: Optional[List[str]] = None
 
 
 class RuntimeParameters(BaseModel):
-    properties: Dict[str, str] = {}
+    properties: Dict[str, Any] = {}
 
 
 class SearchCustomization(BaseModel):
-    featuresFilter: Optional[FeaturesFilter]
-    extractFeatures: Optional[bool]
-    accurateModel: Optional[bool]
-    importanceThreshold: Optional[float]
-    maxFeatures: Optional[int]
-    returnScores: Optional[bool]
-    runtimeParameters: Optional[RuntimeParameters]
-    metricsCalculation: Optional[bool]
+    featuresFilter: Optional[FeaturesFilter] = None
+    extractFeatures: Optional[bool] = None
+    accurateModel: Optional[bool] = None
+    importanceThreshold: Optional[float] = None
+    maxFeatures: Optional[int] = None
+    returnScores: Optional[bool] = None
+    runtimeParameters: Optional[RuntimeParameters] = None
+    metricsCalculation: Optional[bool] = None
 
     def __repr__(self):
         return (

upgini/normalizer/normalize_utils.py

@@ -0,0 +1,202 @@
+import hashlib
+from logging import Logger, getLogger
+from typing import Dict, List
+
+import numpy as np
+import pandas as pd
+from pandas.api.types import is_bool_dtype as is_bool
+from pandas.api.types import is_datetime64_any_dtype as is_datetime
+from pandas.api.types import (
+    is_float_dtype,
+    is_numeric_dtype,
+    is_object_dtype,
+    is_string_dtype,
+)
+
+from upgini.errors import ValidationError
+from upgini.metadata import (
+    ENTITY_SYSTEM_RECORD_ID,
+    EVAL_SET_INDEX,
+    SEARCH_KEY_UNNEST,
+    SYSTEM_RECORD_ID,
+    TARGET,
+    SearchKey,
+)
+from upgini.resource_bundle import ResourceBundle, get_custom_bundle
+from upgini.utils import find_numbers_with_decimal_comma
+from upgini.utils.datetime_utils import DateTimeSearchKeyConverter
+from upgini.utils.phone_utils import PhoneSearchKeyConverter
+from upgini.utils.warning_counter import WarningCounter
+
+
+class Normalizer:
+
+    MAX_STRING_FEATURE_LENGTH = 24573
+
+    def __init__(
+        self,
+        search_keys: Dict[str, SearchKey],
+        generated_features: List[str],
+        bundle: ResourceBundle = None,
+        logger: Logger = None,
+        warnings_counter: WarningCounter = None,
+        silent_mode=False,
+    ):
+        self.search_keys = search_keys
+        self.generated_features = generated_features
+        self.bundle = bundle or get_custom_bundle()
+        self.logger = logger or getLogger()
+        self.warnings_counter = warnings_counter or WarningCounter()
+        self.silent_mode = silent_mode
+        self.columns_renaming = {}
+
+    def normalize(self, df: pd.DataFrame) -> pd.DataFrame:
+        df = df.copy()
+        df = self._rename_columns(df)
+
+        df = self._remove_dates_from_features(df)
+
+        df = self._cut_too_long_string_values(df)
+
+        df = self._convert_bools(df)
+
+        df = self._convert_float16(df)
+
+        df = self._correct_decimal_comma(df)
+
+        df = self._convert_phone_numbers(df)
+
+        df = self.__convert_features_types(df)
+
+        return df
+
+    def _rename_columns(self, df: pd.DataFrame):
+        # logger.info("Replace restricted symbols in column names")
+        new_columns = []
+        dup_counter = 0
+        for column in df.columns:
+            if column in [
+                TARGET,
+                EVAL_SET_INDEX,
+                SYSTEM_RECORD_ID,
+                ENTITY_SYSTEM_RECORD_ID,
+                SEARCH_KEY_UNNEST,
+                DateTimeSearchKeyConverter.DATETIME_COL,
+            ] + self.generated_features:
+                self.columns_renaming[column] = column
+                new_columns.append(column)
+                continue
+
+            new_column = str(column)
+            suffix = hashlib.sha256(new_column.encode()).hexdigest()[:6]
+            if len(new_column) == 0:
+                raise ValidationError(self.bundle.get("dataset_empty_column_names"))
+            # db limit for column length
+            if len(new_column) > 250:
+                new_column = new_column[:250]
+
+            # make column name unique relative to server features
+            new_column = f"{new_column}_{suffix}"
+
+            new_column = new_column.lower()
+
+            # if column starts with non alphabetic symbol then add "a" to the beginning of string
+            if ord(new_column[0]) not in range(ord("a"), ord("z") + 1):
+                new_column = "a" + new_column
+
+            # replace unsupported characters to "_"
+            for idx, c in enumerate(new_column):
+                if ord(c) not in range(ord("a"), ord("z") + 1) and ord(c) not in range(ord("0"), ord("9") + 1):
+                    new_column = new_column[:idx] + "_" + new_column[idx + 1 :]
+
+            if new_column in new_columns:
+                new_column = f"{new_column}_{dup_counter}"
+                dup_counter += 1
+            new_columns.append(new_column)
+
+            # df.columns.values[col_idx] = new_column
+            # rename(columns={column: new_column}, inplace=True)
+
+            if new_column != column and column in self.search_keys:
+                self.search_keys[new_column] = self.search_keys[column]
+                del self.search_keys[column]
+            self.columns_renaming[new_column] = str(column)
+        df.columns = new_columns
+        return df
+
+    def _get_features(self, df: pd.DataFrame) -> List[str]:
+        system_columns = [ENTITY_SYSTEM_RECORD_ID, EVAL_SET_INDEX, SEARCH_KEY_UNNEST, SYSTEM_RECORD_ID, TARGET]
+        features = set(df.columns) - set(self.search_keys.keys()) - set(system_columns)
+        return sorted(list(features))
+
+    def _remove_dates_from_features(self, df: pd.DataFrame):
+        features = self._get_features(df)
+
+        removed_features = []
+        for f in features:
+            if is_datetime(df[f]) or isinstance(df[f].dtype, pd.PeriodDtype):
+                removed_features.append(f)
+                df.drop(columns=f, inplace=True)
+
+        if removed_features:
+            msg = self.bundle.get("dataset_date_features").format(removed_features)
+            self.logger.warning(msg)
+            if not self.silent_mode:
+                print(msg)
+            self.warnings_counter.increment()
+
+        return df
+
+    def _cut_too_long_string_values(self, df: pd.DataFrame):
+        """Check that string values less than maximum characters for LLM"""
+        # logger.info("Validate too long string values")
+        for col in df.columns:
+            if is_string_dtype(df[col]) or is_object_dtype(df[col]):
+                max_length: int = df[col].astype("str").str.len().max()
+                if max_length > self.MAX_STRING_FEATURE_LENGTH:
+                    df[col] = df[col].astype("str").str.slice(stop=self.MAX_STRING_FEATURE_LENGTH)
+
+        return df
+
+    @staticmethod
+    def _convert_bools(df: pd.DataFrame):
+        """Convert bool columns to string"""
+        # logger.info("Converting bool to int")
+        for col in df.columns:
+            if is_bool(df[col]):
+                df[col] = df[col].astype("str")
+        return df
+
+    @staticmethod
+    def _convert_float16(df: pd.DataFrame):
+        """Convert float16 to float"""
+        # logger.info("Converting float16 to float")
+        for col in df.columns:
+            if is_float_dtype(df[col]):
+                df[col] = df[col].astype("float64")
+        return df
+
+    def _correct_decimal_comma(self, df: pd.DataFrame):
+        """Check DataSet for decimal commas and fix them"""
+        # logger.info("Correct decimal commas")
+        columns_to_fix = find_numbers_with_decimal_comma(df)
+        if len(columns_to_fix) > 0:
+            self.logger.warning(f"Convert strings with decimal comma to float: {columns_to_fix}")
+            for col in columns_to_fix:
+                df[col] = df[col].astype("string").str.replace(",", ".", regex=False).astype(np.float64)
+        return df
+
+    def _convert_phone_numbers(self, df: pd.DataFrame) -> pd.DataFrame:
+        maybe_country_col = SearchKey.find_key(self.search_keys, SearchKey.COUNTRY)
+        for phone_col in SearchKey.find_all_keys(self.search_keys, SearchKey.PHONE):
+            converter = PhoneSearchKeyConverter(phone_col, maybe_country_col)
+            df = converter.convert(df)
+        return df
+
+    def __convert_features_types(self, df: pd.DataFrame):
+        # self.logger.info("Convert features to supported data types")
+
+        for f in self._get_features(df):
+            if not is_numeric_dtype(df[f]):
+                df[f] = df[f].astype("string")
+        return df

upgini/utils/country_utils.py

@@ -4,6 +4,22 @@ from pandas.api.types import is_object_dtype, is_string_dtype
 from upgini.utils.base_search_key_detector import BaseSearchKeyDetector
 
 
+class CountrySearchKeyConverter:
+
+    def __init__(self, country_col: str):
+        self.country_col = country_col
+
+    def convert(self, df: pd.DataFrame) -> pd.DataFrame:
+        df[self.country_col] = (
+            df[self.country_col]
+            .astype("string")
+            .str.upper()
+            .str.replace(r"[^A-Z]", "", regex=True)
+            .str.replace("UK", "GB", regex=False)
+        )
+        return df
+
+
 class CountrySearchKeyDetector(BaseSearchKeyDetector):
     def _is_search_key_by_name(self, column_name: str) -> bool:
         return "country" in str(column_name).lower()