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)