embed_rerank_model.py

from openai import OpenAI
import yaml
import logging
import numpy as np

class BaseModel:
    def __init__(self, model_type, model_name, config_path='./configs/models.yaml', logger=None):
        self.config = self.load_config(config_path)
        self.model_config = self.config[model_type][model_name]
        self.model = self.model_config['model']
        self.logger = logger if logger is not None else self._default_logger(model_type)
        self.local_api_key = self.model_config['local_api_key']
        self.local_base_url = self.model_config['local_base_url']
        self.logger.info(f'[{model_type}] Initializing Model: {model_name}')
        self.client = OpenAI(api_key=self.local_api_key, base_url=self.local_base_url)
    
    def load_config(self, path):
        with open(path, 'r', encoding='utf-8') as file:
            return yaml.safe_load(file)
        
    def _default_logger(self, model_type):
        logger = logging.getLogger(f"{model_type}Logger")
        logger.addHandler(logging.NullHandler())
        return logger

class EmbeddingModel(BaseModel):
    def __init__(self, embedding_model, config_path='./configs/models.yaml', logger=None):
        super().__init__('embedding_models', embedding_model, config_path, logger)
    
    def embed_query(self, query=None):
        response = self.client.embeddings.create(
            input=query,
            model=self.model
        )
        return response.data[0].embedding
    
    def embed_documents(self, documents=None):
        response = self.client.embeddings.create(
            input=documents,
            model=self.model
        )
        embs =  [response.data[i].embedding for i in range(len(response.data))]
        return embs

class RerankingModel(BaseModel):
    def __init__(self, reranking_model, config_path='./llm_tools/configs/models.yaml', logger=None):
        super().__init__('reranking_models', reranking_model, config_path, logger)
    
    def rerank_query(self, input=None, query=None):
        response = self.client.embeddings.create(
            model=self.model,
            input=input,
            extra_body={"query": query}
        )
        return response.data[0].embedding
    
    def rerank_documents(self, documents=None, query=None):
        response = self.client.embeddings.create(
            model=self.model,
            input=documents,
            extra_body={"query": query}
        )
        scores = [response.data[i].embedding for i in range(len(response.data))]
        return scores

# 使用範例
if __name__ == '__main__':
    #### EmbeddingModel ####
    embmodel = EmbeddingModel(embedding_model="m3e-base", config_path='./configs/models.yaml')
    
    # Single query embedding
    input = "The food was delicious and the waiter was friendly."
    query_embedding = np.array(embmodel.embed_query(input))
    print(len(query_embedding))

    # Document embedding
    documents = ["The food was delicious and the waiter was friendly.",
                 "The service was slow and the food was not very good."]
    document_embeddings = np.array(embmodel.embed_documents(documents))
    print(document_embeddings.shape)

    
    #### RerankingModel ####
    rerankmodel = RerankingModel(reranking_model="bge-reranker-large", config_path='./configs/models.yaml')
    
    # Single query reranking
    input = "The food was delicious and the waiter was friendly."
    query = "The food was delicious and the waiter was friendly."
    query_reranking = rerankmodel.rerank_query(input, query)
    print(query_reranking)

    # Document reranking
    documents = ["The food was delicious and the waiter was friendly.",
                 "The service was slow and the food was not very good."]
    query = "The food was delicious and the waiter was friendly."
    document_reranking = np.array(rerankmodel.rerank_documents(documents, query))
    print(document_reranking)