src/Database.py

import pandas as pd
import os
from tqdm import tqdm

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity


class Database:
 def __init__(self, parquet_path=None, customized_parquets = None):
 self.default_parquet_path = 'datas/database_4000.parquet'
 self.parquet_path = parquet_path or self.default_parquet_path

 self.default_customized_parquets = ["datas/customized_database_0.parquet"]
 self.customized_parquets = customized_parquets or self.default_customized_parquets

 self.datas = None
 self.last_save_table = None
 
 if os.path.exists(self.parquet_path):
 self.load_from_parquet(self.parquet_path)

 self.load_from_customized(self.customized_parquets)

 self.clip_extractor = None
 self.bge_extractor = None

 self.en_keyword2data = {}

 def build_en_keyword2index(self):
 # build in lower case
 self.en_keyword2data = {row['translated_word'].lower(): row for i, row in self.datas.iterrows()}

 def search_by_en_keyword(self, keyword):
 if len(self.en_keyword2data) == 0:
 self.build_en_keyword2index()

 keyword = keyword.lower()
 if keyword in self.en_keyword2data:
 ans = self.en_keyword2data[keyword].to_dict()
 del ans["clip_feature"]
 del ans["bge_feature"]
 return ans
 else:
 return None
 
 def load_from_parquet(self, parquet_path):
 self.datas = pd.read_parquet(parquet_path)

 def load_from_customized(self, customized_parquets=None):
 customized_parquets = customized_parquets or self.customized_parquets
 
 # Load each parquet file and concatenate them into the self.datas DataFrame
 for index, parquet_file in enumerate(customized_parquets):
 if os.path.exists(parquet_file):
 temp_df = pd.read_parquet(parquet_file)
 if self.datas is None:
 self.datas = temp_df
 else:
 self.datas = pd.concat([self.datas, temp_df], ignore_index=True)

 # if last parquet file
 if index == len(customized_parquets) - 1:
 self.last_save_table = temp_df

 # if customized_parquets:
 # Record the last parquet file's contents as self.last_save_table
 

 def add_data(self, data, if_save=True):
 required_columns = ['keyword', 'name_in_cultivation', 'description_in_cultivation', 'translated_word', 'description']
 for column in required_columns:
 if column not in data:
 raise ValueError(f"Missing required field: {column}")
 
 # Optional field
 if 'founder' not in data:
 data['founder'] = ""
 
 # Extract features
 if self.clip_extractor is None:
 self.init_clip_extractor()
 if self.bge_extractor is None:
 self.init_bge_extractor()

 data['clip_feature'] = self.clip_extractor.extract_text(data['translated_word'] + '.' + data['description'])
 data['bge_feature'] = self.bge_extractor.extract([data['keyword']])[0].tolist()

 # Convert to DataFrame and add to self.datas
 data_df = pd.DataFrame([data])
 if self.datas is None:
 self.datas = data_df
 else:
 self.datas = pd.concat([self.datas, data_df], ignore_index=True)

 # set self.en_keyword2data to last row of self.datas
 self.en_keyword2data[data['translated_word'].lower()] = self.datas.iloc[-1]
 
 # Add to last_save_table
 if self.last_save_table is None:
 # self.last_save_table = data_df
 # create a new DataFrame with the same columns as self.datas
 self.last_save_table = pd.DataFrame(columns=self.datas.columns)
 
 self.last_save_table = pd.concat([self.last_save_table, data_df], ignore_index=True)
 
 if if_save:
 self.save_to_parquet(self.customized_parquets[-1], self.last_save_table )

 def add_datas(self, datas, if_save=True):
 for data in datas:
 self.add_data(data, if_save=False)
 if if_save:
 self.save_to_parquet(self.customized_parquets[-1], self.last_save_table)

 def init_from_excel(self, excel_path):
 df = pd.read_excel(excel_path)
 
 # Drop rows with any empty cell in the required columns
 df.dropna(subset=['keyword', 'name_in_cultivation', 'description_in_cultivation', 'translated_word', 'description'], inplace=True)
 
 # Add the new columns
 df['clip_feature'] = None
 df['bge_feature'] = None
 
 self.datas = df

 self.extract_clip()
 self.extract_bge()
 
 def save_to_parquet(self, parquet_path=None, df = None):
 
 parquet_path = parquet_path or self.default_parquet_path
 if df is None:
 if self.datas is not None:
 self.datas.to_parquet(parquet_path)
 else:
 df.to_parquet(parquet_path)

 def init_clip_extractor(self):
 if self.clip_extractor is None:
 try:
 from CLIPExtractor import CLIPExtractor
 except:
 from src.CLIPExtractor import CLIPExtractor

 cache_dir = "models"

 self.clip_extractor = CLIPExtractor(model_name = "openai/clip-vit-large-patch14",cache_dir = cache_dir)

 
 def extract_clip(self):
 if self.clip_extractor is None:
 self.init_clip_extractor()

 clip_features = []
 # for text in tqdm(self.datas['keyword'], desc='Extracting CLIP features'):
 for index, row in tqdm(self.datas.iterrows(), desc='Extracting CLIP features', total=len(self.datas)):
 text = row['translated_word'] + '.' + row['description']
 if text:
 feature = self.clip_extractor.extract_text(text)
 else:
 feature = None
 clip_features.append(feature)

 self.datas['clip_feature'] = clip_features

 def init_bge_extractor(self):
 if self.bge_extractor is None:
 try:
 from text_embedding import TextExtractor
 except:
 from src.text_embedding import TextExtractor

 self.bge_extractor = TextExtractor('BAAI/bge-small-zh-v1.5')

 def top_k_search(self, query_feature, attribute, top_k=15):
 # Ensure the attribute exists in the dataframe
 if attribute not in self.datas.columns:
 raise ValueError(f"Attribute {attribute} not found in the data.")
 
 # Convert query feature and attribute features to numpy arrays
 query_feature = np.array(query_feature).reshape(1, -1)
 attribute_features = np.stack(self.datas[attribute].dropna().values)
 
 # Compute cosine similarity between query and all attributes
 similarities = cosine_similarity(query_feature, attribute_features)[0]
 
 # Get the top_k indices based on similarity
 top_k_indices = np.argsort(similarities)[-top_k:][::-1]
 
 # Retrieve the top_k most similar items
 top_k_results = self.datas.iloc[top_k_indices].copy()
 
 top_k_results = top_k_results.drop(columns=['clip_feature', 'bge_feature'])

 top_k_results['similarity'] = similarities[top_k_indices]

 return top_k_results.to_dict(orient='records')

 def search_with_image_name(self, image_name):
 self.init_clip_extractor()

 img_feature = self.clip_extractor.extract_image_from_file(image_name)

 return self.top_k_search(img_feature, 'clip_feature')

 def search_with_image(self, image, if_opencv = False ):
 if self.clip_extractor is None:
 self.init_clip_extractor()

 img_feature = self.clip_extractor.extract_image(image, if_opencv = if_opencv)

 return self.top_k_search(img_feature, 'clip_feature')
 
 def search_with_chinese(self, text):
 if self.bge_extractor is None:
 self.init_bge_extractor()

 text_feature = self.bge_extractor.extract([text])[0].tolist()

 return self.top_k_search(text_feature, 'bge_feature')



 def extract_bge(self):
 if self.bge_extractor is None:
 self.init_bge_extractor()
 
 # Extract features for each row and store them in the bge_feature column
 bge_features = []
 for text in tqdm(self.datas['keyword'], desc='Extracting BGE features'):
 if text:
 feature = self.bge_extractor.extract([text])[0].tolist()
 else:
 feature = None
 bge_features.append(feature)
 
 self.datas['bge_feature'] = bge_features

if __name__ == '__main__':
 # Usage example
 db = Database()
 re_extract = False
 if db.datas is None or re_extract:
 print("Rebuilding database from excel file")
 db.init_from_excel('datas/database_4000.xlsx')
 db.save_to_parquet()

 # print(db.datas[0].keys())

 query_text = "钢琴"

 results = db.search_with_chinese(query_text)

 print(results[0].keys())

 for result in results[:3]:
 print(result)

 image_path = "datas/老虎.jpg"

 results = db.search_with_image_name(image_path)

 for result in results[:3]:
 print(result)
# 'keyword': '老虎狗', 'name_in_cultivation': '灵虎犬神', 'description_in_cultivation': '在九天灵脉汇聚的仙山之巅，灵虎犬神身披星图
# 斑纹，汲取日月精华，以雷霆之力守护仙脉，其双眼中映照着轮回之道，是修仙者追寻天地真理的指引，也是象征极致灵性的神秘灵兽。', 'translated_word': 'Tiger Dog', 'description': 'A Tiger Dog is a term that might refer to a mythical creature or a breed of dog with a distinctive and unusual appearance, resembling the features of a tiger. It could be characterized by its striking coat with patterns similar to those of a tiger, or by having a demeanor that is fierce and majestic like a tiger. This term is not commonly used in 
# conventional contexts and might be found in stories, folktales, or in the names of unique dog breeds that have been bred to exhibit such features.', 'founder': ''
 # test_new_data = {
 # "keyword": "老虎狗2",
 # "name_in_cultivation": "灵虎犬神",
 # "description_in_cultivation": "在九天灵脉汇聚的仙山之巅，灵虎犬神身披星图斑纹，汲取日月精华，以雷霆之力守护仙脉，其双眼中映照着轮回之道，是修仙者追寻天地真理的指引，也是象征极致灵性的神秘灵兽。",
 # "translated_word": "Tiger Dog",
 # "description":"A Tiger Dog is a term that might refer to a mythical creature or a breed of dog with a distinctive and unusual appearance, resembling the features of a tiger. It could be characterized by its striking coat with patterns similar to those of a tiger, or by having a demeanor that is fierce and majestic like a tiger. This term is not commonly used in conventional contexts and might be found in stories, folktales, or in the names of unique dog breeds that have been bred to exhibit such features."
 # }

 # db.add_data(test_new_data)