住所データをダウンロードする機能を追加し、環境変数からターゲットディレクトリを取得するように修正。新たに必要なディレクトリを.gitignoreに追加し、一時ディレクトリの管理を強化。

.gitignore

@@ -2,4 +2,6 @@
 embeddings/
 embeddings_/
 __pycache__/
-.env
+.env
+data/
+digital_agency/

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import zipfile
+from tqdm import tqdm
 from sklearn.metrics.pairwise import cosine_similarity
 from pathlib import Path
 import spacy
@@ -21,7 +23,8 @@ load_dotenv()
 # =========================
 #  Global variables
 # =========================
-# 環境変数からHUGGING_FACE_TOKENを取得
+TARGET_DIR = Path(os.environ.get('TARGET_DIR'))
+
 HUGGING_FACE_TOKEN = os.environ.get('HUGGING_FACE_TOKEN')
 EMBEDDING_MODEL_ENDPOINT = os.environ.get('EMBEDDING_MODEL_ENDPOINT')
 ABRG_ENDPOINT = os.environ.get('ABRG_ENDPOINT')
@@ -34,6 +37,71 @@ VECTOR_SEARCH_COLLECTION_NAME_V2 = os.environ.get('VECTOR_SEARCH_COLLECTION_NAME
 MILVUS_CLIENT = MilvusClient(uri=VECTOR_SEARCH_ENDPOINT, token=VECTOR_SEARCH_TOKEN)
 print(f"Connected to DB: {VECTOR_SEARCH_ENDPOINT} successfully")
 
+# ----------------------------
+#  Download mt_city_all.csv
+# ----------------------------
+temp_dir = Path('temp')
+temp_dir.mkdir(exist_ok=True)
+
+city_all_url = 'https://catalog.registries.digital.go.jp/rsc/address/mt_city_all.csv.zip'
+zip_file_path = temp_dir / 'mt_city_all.csv.zip'
+
+# すでにファイルが存在する場合はダウンロードをスキップ
+if not os.path.exists(zip_file_path):
+    # ZIPファイルをダウンロード
+    response = requests.get(city_all_url)
+    with open(zip_file_path, 'wb') as f:
+        f.write(response.content)
+
+    # target_dir直下に解凍
+    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+        zip_ref.extractall(TARGET_DIR)
+
+# ------------------------------------
+#  Download mt_parcel_cityXXXXXX.csv
+# ------------------------------------
+city_all_file = TARGET_DIR / 'mt_city_all.csv'
+city_all_df = pd.read_csv(city_all_file)
+lg_codes = city_all_df['lg_code'].tolist()
+print('lg_codes', len(lg_codes))
+
+for lg_code in tqdm(lg_codes):
+    parcel_url = f'https://catalog.registries.digital.go.jp/rsc/address/mt_parcel_city{lg_code:06d}.csv.zip'
+    zip_file_path = temp_dir / f'mt_parcel_city{lg_code:06d}.csv.zip'
+
+    if not os.path.exists(TARGET_DIR / 'parcel' / f'mt_parcel_city{lg_code:06d}.csv'):
+        response = requests.get(parcel_url)
+        if response.status_code == 200:  # URLが存在する場合のみ処理を続ける
+            with open(zip_file_path, 'wb') as f:
+                f.write(response.content)
+            with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+                zip_ref.extractall(TARGET_DIR / 'parcel')
+        time.sleep(0.2)  # ダウンロードごとに200msのスリープを入れる
+
+# ------------------------------------
+#  Download mt_rsdtdsp_rsdt_prefXX.csv
+# ------------------------------------
+pref_codes = list(set([('%06d' % lg_code)[0:2] for lg_code in lg_codes]))
+print('pref_codes', len(pref_codes))
+
+for pref_code in tqdm(pref_codes):
+    rsdt_url = f'https://catalog.registries.digital.go.jp/rsc/address/mt_rsdtdsp_rsdt_pref{pref_code}.csv.zip'
+    zip_file_path = temp_dir / f'mt_rsdtdsp_rsdt_pref{pref_code}.csv.zip'
+
+    if not os.path.exists(TARGET_DIR / 'rsdt' / f'mt_rsdtdsp_rsdt_pref{pref_code}.csv.zip'):
+        response = requests.get(parcel_url)
+        if response.status_code == 200:  # URLが存在する場合のみ処理を続ける
+            with open(zip_file_path, 'wb') as f:
+                f.write(response.content)
+            with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+                zip_ref.extractall(TARGET_DIR / 'rsdt')
+        time.sleep(0.2)  # ダウンロードごとに200msのスリープを入れる
+
+# 一時ディレクトリを削除
+for file in temp_dir.iterdir():
+    file.unlink()
+temp_dir.rmdir()
+
 # =========================
 #  Utilitiy functions
 # =========================
@@ -392,8 +460,7 @@ def create_vector_search_tab():
                 }
                 result_df = pd.DataFrame([splits.values()], columns=splits.keys())
             with measure('load city_all_file'):
-                target_dir = Path(r'C:\Users\taish\Development\whispercustom\projects\abr-geocoder\temp\download')
-                city_all_file = target_dir / 'mt_city_all.csv'
+                city_all_file = TARGET_DIR / 'mt_city_all.csv'
                 city_all_df = pd.read_csv(city_all_file)
                 city_all_df_temp = city_all_df[city_all_df['pref'] == splits['pref']]
                 city_name1 = city_all_df_temp['county'].fillna('') + city_all_df_temp['city'].fillna('') + city_all_df_temp['ward'].fillna('')
@@ -403,7 +470,7 @@ def create_vector_search_tab():
                     raise Exception('Too many lg_code')
                 lg_code = lg_codes[0]
             with measure('load parcel_city_file'):
-                parcel_city_file = target_dir / f'mt_parcel_city{lg_code:06d}.csv'
+                parcel_city_file = TARGET_DIR / 'parcel' / f'mt_parcel_city{lg_code:06d}.csv'
                 if not os.path.exists(parcel_city_file):
                     # raise gr.Error('Too many lg_code')
                     raise Exception('Too many lg_code')
@@ -446,8 +513,7 @@ def create_vector_search_tab():
                 ]
             with measure('load rsdtdsp_file'):
                 pref_code = ('%06d' % lg_code)[0:2]
-                rsdtdsp_dir = Path(rf'G:\マイドライブ\Development\Dataset\Misc\japanese_address\rsdt\original')
-                rsdtdsp_file = rsdtdsp_dir / f'mt_rsdtdsp_rsdt_pref{pref_code}.csv\mt_rsdtdsp_rsdt_pref{pref_code}.csv'
+                rsdtdsp_file = RSDTDSP_DIR / f'mt_rsdtdsp_rsdt_pref{pref_code}.csv'
                 if not os.path.exists(rsdtdsp_file):
                     # raise gr.Error(f'Not found: {rsdtdsp_file}')
                     raise Exception(f'Not found: {rsdtdsp_file}')