import pandas as pd
import numpy as np
import os
import sys
str_i = str(sys.argv[1])
npy_file = "img_emb_" + str_i + '.npy'
metadata_file = "metadata_" + str_i + '.parquet'
text_npy = "text_emb_" + str_i + '.npy'
# 모든 파일 로드
im_emb = np.load(npy_file)
text_emb = np.load(text_npy)
data = pd.read_parquet(metadata_file)
# 파일 병합
data = pd.concat([data, pd.DataFrame({"image_embedding" : [*im_emb]}), pd.DataFrame({"text_embedding" : [*text_emb]})], axis=1, copy=False)
# ClickHouse에 임포트할 컬럼
data = data[['url', 'caption', 'NSFW', 'similarity', "image_embedding", "text_embedding"]]
# np.array를 리스트로 변환
data['image_embedding'] = data['image_embedding'].apply(lambda x: x.tolist())
data['text_embedding'] = data['text_embedding'].apply(lambda x: x.tolist())
# caption에 다양한 종류의 따옴표가 포함될 수 있으므로 이 처리가 필요함
data['caption'] = data['caption'].apply(lambda x: x.replace("'", " ").replace('"', " "))
# 데이터를 CSV 파일로 내보내기
data.to_csv(str_i + '.csv', header=False)
# 원본 데이터 파일 삭제
os.system(f"rm {npy_file} {metadata_file} {text_npy}")