# prepare_municipal_embeddings.py # Updated to download municipal HTML, citation, and embeddings Parquet files from laion/american_municipal_law/american_law/data, process them, and create CID maps import os import logging import requests from datasets import load_dataset, concatenate_datasets import pickle logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) MUNICIPAL_HTML_PATH = "/data/municipal_html" MUNICIPAL_CITATION_PATH = "/data/municipal_citation" MUNICIPAL_EMBEDDINGS_PATH = "/data/municipal_embeddings" MUNICIPAL_CID_TO_HTML_PATH = "/data/cid_to_html.pkl" MUNICIPAL_CID_TO_CITATION_PATH = "/data/cid_to_citation.pkl" base_url = "https://huggingface.co/datasets/laion/american_municipal_law/resolve/main/american_law/data" def download_file_if_not_exists(url, save_path): if not os.path.exists(save_path): try: response = requests.get(url) response.raise_for_status() with open(save_path, 'wb') as f: f.write(response.content) logger.info(f"Downloaded and saved file to {save_path}") except Exception as e: logger.error(f"Failed to download from {url}: {str(e)}") def main(): try: logger.info("Starting prepare_municipal_embeddings.py") # List of unique GNIS IDs extracted from the "parquet names.docx" document gnis_list = [ gnis_list = [ "1008538", "1008539", "1008540", "1008543", "1008547", "1008550", "1008552", "1008559", "1008564", "1008566", "1026123", "1026124", "1026125", "1026128", "1026328", "1026329", "1026336", "1026337", "1026340", "1026341", "1035989", "1036064", "1036100", "1036131", "1036146", "1036165", "1036231", "1036246", "1036306", "1036307", "1036311", "1036317", "1085949", "1085956", "1085987", "1086167", "1086201", "1086669", "1086677", "1215071", "1215263", "1215416", "1220064", "1220067", "1220069", "1220071", "1244251", "1244717", "1244832", "1245666", "1247979", "1247981", "1247982", "1247985", "1247987", "1247988", "1247991", "1247992", "1247993", "1247998", "1247999", "1248000", "1248001", "1248002", "1248003", "1248005", "1248006", "1248007", "1248008", "1248009", "1267477", "1267547", "1267560", "1267567", "1388537", "1462060", "1462212", "1480147", "1494590", "1496656", "1496705", "1497748", "1498359", "1498416", "1498417", "1498422", "1498423", "1498555", "1498556", "1498557", "1498559", "1498563", "1500545", "1500747", "1500936", "1501122", "1513272", "1529225", "1529226", "1531819", "1531822", "1531823", "1533501", "1569014", "156909", "1581066", "1581071", "1581073", "1581087", "1581088", "1581094", "1581096", "1581097", "1581100", "1581104", "1581108", "1581114", "1582777", "1582791", "1582958", "1583034", "1583036", "1583076", "1583269", "1583281", "1583462", "1583492", "1583527", "1583581", "1583593", "1583619", "1583625", "1583673", "1583695", "1583699", "1583701", "1583711", "1583728", "1583757", "1583779", "1583784", "1583829", "1583859", "1583885", "1583890", "1583892", "1583913", "1583915", "1583916", "1583929", "1583947", "1583953", "1583991", "1583997", "1584014", "1584039", "1584041", "1584045", "1584056", "1584092", "1584112", "1584133", "1584157", "1584181", "1584187", "1584214", "1584233", "1584237", "1584239", "1584250", "1584255", "1584273", "1584282", "1584328", "1584374", "1584377", "1584427", "1622945", "1623022", "1625803", "1625813", "1625818", "1625826", "1625953", "1625965", "1626017", "1626049", "1626052", "1626085", "1626089", "1626132", "1626165", "1626181", "1626206", "1626207", "1626330", "1626369", "1626411", "1626558", "1626561", "1626646", "1626672", "1626699", "1626771", "1626821", "1627124", "1627209", "1669605", "1673664", "1682927", "1687999", "1688005", "1688166", "1789068", "1789070", "198150", "198160", "198174", "2378275", "2378292", "2390236", "2390580", "2390581", "2390591", "2390622", "2390723", "2390767", "2390787", "2391326", "2391405", "2391470", "2393515", "2394023", "2394038", "2394489", "2394586", "2394604", "2394614", "2394706", "2394924", "2395032", "2395201", "2395731", "2395846", "2395865", "2396049", "2396415", "2396420", "2396453", "2396517", "2396526", "2396534", "2396567", "2396647", "2396662", "2396754", "2396794", "2396821", "2396945", "2396972", "2396979", "2397028", "2397049", "2397097", "2397106", "2397111", "2397119", "2397155", "2397159", "2397168", "2397174", "2397183", "2397195", "2397199", "2397211", "2397222", "2397235", "2397241", "2397242", "2397257", "2397264", "2397267", "2397268", "2397273", "2397294", "2397304", "2397323", "2397327", "2397328", "2397334", "2397369", "2397392", "2397458", "2397465", "2397467", "2397489", "2397490", "2397502", "2397517", "2397558", "2397715", "2397990", "2398512", "2398518", "2398647", "2398701", "2398754", "2398819", "2398935", "2398968", "2398972", "2399184", "2399218", "2399387", "2399789", "2399970", "2400087", "2400108", "2400148", "2400153", "2400161", "2400163", "2403060", "2403064", "2403077", "2403079", "2403080", "2403087", "2403093", "2403097", "2403098", "2403100", "2403103", "2403104", "2403105", "2403106", "2403107", "2403109", "2403123", "2403126", "2403127", "2403128", "2403129", "2403132", "2403134", "2403135", "2403139", "2403141", "2403143", "2403144", "2403146", "2403147", "2403149", "2403151", "2403163", "2403514", "2403516", "2403518", "2403519", "2403521", "2403536", "2403537", "2403541", "2403543", "2403544", "2403549", "2403558", "2403559", "2403563", "2403567", "2403592", "2403617", "2403633", "2403635", "2403649", "2403667", "2403674", "2403693", "2403703", "2403729", "2403742", "2403758", "2403762", "2403785", "2403794", "2403812", "2403816", "2403825", "2403872", "2403875", "2403881", "2403887", "2403908", "2403937", "2403945", "2403946", "2403949", "2403988", "2403989", "2403990", "2403993", "2403998", "2404005", "2404011", "2404019", "2404044", "2404079", "2404093", "2404096", "2404150", "2404219", "2404226", "2404232", "2404234", "2404236", "2404237", "2404238", "2404239", "2404244", "2404245", "2404249", "2404250", "2404258", "2404261", "2404262", "2404263", "2404269", "2404301", "2404347", "2404352", "2404354", "2404355", "2404357", "2404361", "2404365", "2404366", "2404367", "2404369", "2404370", "2404371", "2404376", "2404377", "2404381", "2404384", "2404389", "2404391", "2404392", "2404393", "2404395", "2404397", "2404400", "2404407", "2404408", "2404413", "2404414", "2404416", "2404417", "2404424", "2404428", "2404429", "2404433", "2404434", "2404435", "2404437", "2404439", "2404440", "2404443", "2404448", "2404451", "2404453", "2404454", "2404456", "2404458", "2404459", "2404461", "2404464", "2404466", "2404469", "2404471", "2404479", "2404482", "2404485", "2404486", "2404489", "2404495", "2404496", "2404497", "2404499", "2404503", "2404504", "2404505", "2404510", "2404513", "2404514", "2404521", "2404534", "2404536", "2404537", "2404538", "2404548", "2404550", "2404552", "2404553", "2404559", "2404567", "2404568", "2404581", "2404583", "2404587", "2404588", "2404590", "2404592", "2404597", "2404600", "2404604", "2404605", "2404610", "2404615", "2404616", "2404619", "2404624", "2404626", "2404630", "2404632", "2404635", "2404639", "2404643", "2404644", "2404648", "2404650", "2404657", "2404660", "2404661", "2404665", "2404666", "2404675", "2404681", "2404683", "2404688", "2404691", "2404756", "2404758", "2404760", "2404762", "2404765", "2404771", "2404790", "2404808", "2404843", "2404856", "2404955", "2404983", "2404984", "2405009", "2405037", "2405039", "2405040", "2405043", "2405047", "2405061", "2405072", "2405073", "2405074", "2405076", "2405119", "2405122", "2405145", "2405148", "2405152", "2405160", "2405170", "2405177", "2405180", "2405203", "2405204", "2405233", "2405235", "2405245", "2405251", "2405260", "2405280", "2405291", "2405314", "2405318", "2405338", "2405345", "2405351", "2405376", "2405381", "2405413", "2405415", "2405418", "2405432", "2405437", "2405480", "2405525", "2405526", "2405527", "2405543", "2405546", "2405548", "2405550", "2405551", "2405554", "2405558", "2405559", "2405560", "2405564", "2405575", "2405599", "2405600", "2405634", "2405641", "2405691", "2405697", "2405699", "2405742", "2405767", "2405802", "2405804", "2405857", "2405882", "2405904", "2405918", "2405939", "2406060", "2406113", "2406115", "2406132", "2406153", "2406170", "2406246", "2406261", "2406269", "2406320", "2406482", "2406500", "2406520", "2406565", "2406578", "2406583", "2406599", "2406602", "2406608", "2406609", "2406614", "2406616", "2406625", "2406627", "2406668", "2406693", "2406732", "2406749", "2406821", "2406844", "2406845", "2406850", "2406876", "2406923", "2406930", "2407065", "2407290", "2407359", "2407364", "2407367", "2407449", "2407456", "2407535", "2407540", "2407569", "2409400", "2409451", "2409496", "2409679", "2409687", "2409688", "2409689", "2409692", "2409694", "2409801", "2409843", "2409849", "2409850", "2409854", "2409855", "2409856", "2409865", "2409872", "2409991", "2410155", "2410175", "2410183", "2410210", "2410383", "2410491", "2410493", "2410504", "2410510", "2410513", "2410527", "2410528", "2410529", "2410543", "2410594", "2410628", "2410673", "2410764", "2410807", "2410851", "2410999", "2411080", "2411118", "2411175", "2411214", "2411215", "2411244", "2411254", "2411282", "2411354", "2411360", "2411400", "2411427", "2411533", "2411726", "2411742", "2411753", "2411773", "2411883", "2411885", "2411886", "2411887", "2411889", "2411892", "2411897", "2411900", "2411901", "2411904", "2411906", "2411920", "2411921", "2411924", "2411930", "2411931", "2411936", "2411941", "2411949", "2411952", "2411954", "2411960", "2411965", "2411967", "2411976", "2411978", "2411982", "2411995", "2412015", "2412016", "2412019", "2412023", "2412027", "2412036", "2412040", "2412042", "2412048", "2412054", "2412058", "2412064", "2412078", "2412080", "2412088", "2412090", "2412101", "2412102", "2412105", "2412110", "2412117", "2412247", "2412278", "2412308", "2412456", "2412551", "2412573", "2412683", "2412701", "2412708", "2412710", "2412716", "2412740", "2412749", "2412765", "2412770", "2412776", "2412807", "2412808", "2412820", "2412823", "2412828", "2412838", "2412857", "2412863", "2412871", "2412878", "2412900", "2412903", "2412905", "2412910", "2412917", "2412922", "2412930", "2412936", "2412939", "2412970", "2412976", "2412977", "2412992", "2413000", "2413003", "2413004", "2413005", "2413009", "2413013", "2413017", "2413019", "2413026", "2413037", "2413042", "2413053", "2413058", "2413061", "2413088", "2413122", "2413132", "2413286", "2413324", "2413480", "2413501", "2413509", "2413561", "2413593", "277294", "295733", "295752", "306915", "306918", "306921", "307624", "307626", "308547", "308548", "308549", "326638", ] html_datasets = [] citation_datasets = [] embeddings_datasets = [] for gnis in gnis_list: html_url = f"{base_url}/{gnis}_html.parquet?download=true" citation_url = f"{base_url}/{gnis}_citation.parquet?download=true" embeddings_url = f"{base_url}/{gnis}_embeddings.parquet?download=true" html_path = f"/data/{gnis}_html.parquet" citation_path = f"/data/{gnis}_citation.parquet" embeddings_path = f"/data/{gnis}_embeddings.parquet" download_file_if_not_exists(html_url, html_path) download_file_if_not_exists(citation_url, citation_path) download_file_if_not_exists(embeddings_url, embeddings_path) # Load individual datasets html_dataset = load_dataset("parquet", data_files={"train": html_path})["train"] citation_dataset = load_dataset("parquet", data_files={"train": citation_path})["train"] embeddings_dataset = load_dataset("parquet", data_files={"train": embeddings_path})["train"] html_datasets.append(html_dataset) citation_datasets.append(citation_dataset) embeddings_datasets.append(embeddings_dataset) # Combine datasets combined_html = concatenate_datasets(html_datasets) combined_citation = concatenate_datasets(citation_datasets) combined_embeddings = concatenate_datasets(embeddings_datasets) logger.info(f"Combined datasets: HTML ({len(combined_html)}), Citation ({len(combined_citation)}), Embeddings ({len(combined_embeddings)})") # Save combined datasets to disk combined_html.save_to_disk(MUNICIPAL_HTML_PATH) logger.info(f"Saved HTML dataset to {MUNICIPAL_HTML_PATH}") combined_citation.save_to_disk(MUNICIPAL_CITATION_PATH) logger.info(f"Saved citation dataset to {MUNICIPAL_CITATION_PATH}") combined_embeddings.save_to_disk(MUNICIPAL_EMBEDDINGS_PATH) logger.info(f"Saved embeddings dataset to {MUNICIPAL_EMBEDDINGS_PATH}") # Create CID to HTML map cid_to_html = {row['cid']: row['html'] for row in combined_html} with open(MUNICIPAL_CID_TO_HTML_PATH, 'wb') as f: pickle.dump(cid_to_html, f) logger.info(f"Saved CID to HTML map to {MUNICIPAL_CID_TO_HTML_PATH}") # Create CID to citation map cid_to_citation = {row['cid']: row['bluebook_citation'] for row in combined_citation} with open(MUNICIPAL_CID_TO_CITATION_PATH, 'wb') as f: pickle.dump(cid_to_citation, f) logger.info(f"Saved CID to citation map to {MUNICIPAL_CID_TO_CITATION_PATH}") # Download all metadata JSON files metadata_base_url = "https://huggingface.co/datasets/laion/american_municipal_law/resolve/main/american_law/metadata" metadata_dir = "/data/metadata" os.makedirs(metadata_dir, exist_ok=True) for gnis in gnis_list: json_url = f"{metadata_base_url}/{gnis}.json?download=true" json_path = f"{metadata_dir}/{gnis}.json" download_file_if_not_exists(json_url, json_path) logger.info("Downloaded all metadata JSON files") except Exception as e: logger.error(f"Error preparing municipal embeddings: {str(e)}", exc_info=True) if __name__ == "__main__": main()