update files(6)

src/post_extraction_tools/data_quality_enhancer.py

@@ -54,6 +54,12 @@ key_industry_types = [
     "Aviation & Airlines"
 ]
 
+def clean_company_name(name: str) -> str:
+    name = name.lower()
+    name = re.sub(r'\b(inc\.?|ltd\.?|llc\.?|co\.?|corp\.?|corporation|limited|pvt|private|public|plc|gmbh|s\.a\.|srl|bv|ag|oy|ab|spa|sas|sdn bhd|holdings|group|company|enterprises|technologies)\b', '', name)
+    name = re.sub(r'[^a-z0-9\s]', '', name)
+    return [re.sub(r'\s+', ' ', name).strip()]
+
 def enhancer(data: object, embedder) -> list:
     """
     Enhances the data quality by removing duplicates
@@ -73,7 +79,7 @@ def enhancer(data: object, embedder) -> list:
         c1 = companies[i]
         # name1 = tokenize(c1.get("company_name", ""))
         # ind1 = tokenize(c1.get("industry_type", ""))
-        c1_name_embedding = embedder.encode([c1.get("company_name", "")])
+        c1_name_embedding = embedder.encode(clean_company_name(c1.get("company_name", "")))
         c1_ind_embedding = embedder.encode([c1.get("industry_type", "")])
         # c1["ind_embedding"] = c1_ind_embedding
         c1_country = c1.get("country", "").lower().strip()
@@ -83,7 +89,7 @@ def enhancer(data: object, embedder) -> list:
             if j in duplicate_idx:
                 continue
             c2 = companies[j]
-            c2_name_embedding = embedder.encode([c2.get("company_name", "")])
+            c2_name_embedding = embedder.encode(clean_company_name(c2.get("company_name", "")))
             c2_ind_embedding = embedder.encode([c2.get("industry_type", "")])
             # c2["ind_embedding"] = c2_ind_embedding
             c2_country = c2.get("country", "").lower().strip()