| import sys | |
| sys.path.append("../") | |
| from transformers import AutoModel, AutoTokenizer | |
| from transformers.utils import hub | |
| from functools import partial | |
| from bw_utils import get_child_folders | |
| import torch | |
| import os | |
| class EmbeddingModel: | |
| def __init__(self, model_name, language='en'): | |
| self.model_name = model_name | |
| self.language = language | |
| cache_dir = hub.default_cache_path | |
| model_provider = model_name.split("/")[0] | |
| model_smallname = model_name.split("/")[1] | |
| model_path = os.path.join(cache_dir, f"models--{model_provider}--{model_smallname}/snapshots/") | |
| if os.path.exists(model_path) and get_child_folders(model_path): | |
| try: | |
| model_path = os.path.join(model_path,get_child_folders(model_path)[0]) | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| self.model = AutoModel.from_pretrained(model_path) | |
| except Exception as e: | |
| print(e) | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| self.model = AutoModel.from_pretrained(model_name) | |
| else: | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| self.model = AutoModel.from_pretrained(model_name) | |
| def __call__(self, input): | |
| inputs = self.tokenizer(input, return_tensors="pt", padding=True, truncation=True, max_length=256) | |
| with torch.no_grad(): | |
| outputs = self.model(**inputs) | |
| embeddings = outputs.last_hidden_state[:, 0, :].tolist() | |
| return embeddings | |
| class OpenAIEmbedding: | |
| def __init__(self, model_name="text-embedding-ada-002"): | |
| from openai import OpenAI | |
| self.client = OpenAI() | |
| self.model_name = model_name | |
| def __call__(self, input): | |
| if isinstance(input, str): | |
| input = input.replace("\n", " ") | |
| return self.client.embeddings.create(input=[input], model=self.model_name).data[0].embedding | |
| elif isinstance(input,list): | |
| return [self.client.embeddings.create(input=[sentence.replace("\n", " ")], model=self.model_name).data[0].embedding for sentence in input] | |
| def get_embedding_model(embed_name, language='en'): | |
| model_name_dict = { | |
| "bge-m3":"BAAI/bge-m3", | |
| "bge": "BAAI/bge-large-", | |
| "luotuo": "silk-road/luotuo-bert-medium", | |
| "bert": "google-bert/bert-base-multilingual-cased", | |
| } | |
| if embed_name in model_name_dict: | |
| model_name = model_name_dict[embed_name] | |
| if embed_name == 'bge': | |
| model_name += language | |
| return EmbeddingModel(model_name) | |
| else: | |
| return OpenAIEmbedding() | |