Spaces:

silk-road
/

Zero-Haruhi-50_Novels-Playground-API

Running

App Files Files Community

Zero-Haruhi-50_Novels-Playground-API / ChatHaruhi /embeddings.py

silk-road

Upload 15 files

d319ff8 verified 4 months ago

raw history blame contribute delete

No virus

9.11 kB

	import random

	# elif embedding == 'bge_en':
	# embed_name = 'bge_en_s15'
	# elif embedding == 'bge_zh':
	# embed_name = 'bge_zh_s15'

	import torch

	device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


	embedshortname2model_name = {
	"bge_zh":"BAAI/bge-small-zh-v1.5",
	}

	embedname2columnname = {
	"luotuo_openai":"luotuo_openai",
	"openai":"luotuo_openai",
	"bge_zh":"bge_zh_s15",
	"bge_en":"bge_en_s15",
	"bce":"bce_base",
	}

	# 这是用来调试的foo embedding

	def foo_embedding( text ):
	# whatever text input , output a 2 dim 0-1 random vects
	return [random.random(), random.random()]

	# TODO: add bge-zh-small(or family) BCE and openai embedding here 米唯实
	# ======== add bge_zh mmodel
	# by Weishi MI

	def foo_bge_zh_15( text ):
	dim = 512
	model_name = "BAAI/bge-small-zh-v1.5"
	if isinstance(text, str):
	text_list = [text]
	else:
	get_general_embeddings_safe(text, model_name)

	global _model_pool
	global _tokenizer_pool

	if model_name not in _model_pool:
	from transformers import AutoTokenizer, AutoModel
	_tokenizer_pool[model_name] = AutoTokenizer.from_pretrained(model_name)
	_model_pool[model_name] = AutoModel.from_pretrained(model_name)

	_model_pool[model_name].eval()

	# Tokenize sentences
	encoded_input = _tokenizer_pool[model_name](text_list, padding=True, truncation=True, return_tensors='pt', max_length = 512)

	# Compute token embeddings
	with torch.no_grad():
	model_output = _model_pool[model_name](**encoded_input)
	# Perform pooling. In this case, cls pooling.
	sentence_embeddings = model_output[0][:, 0]

	# normalize embeddings
	sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
	return sentence_embeddings.cpu().tolist()[0]
	# return [random.random() for _ in range(dim)]

	def foo_bce( text ):
	from transformers import AutoModel, AutoTokenizer
	if isinstance(text, str):
	text_list = [text]

	# init model and tokenizer
	tokenizer = AutoTokenizer.from_pretrained('maidalun1020/bce-embedding-base_v1')
	model = AutoModel.from_pretrained('maidalun1020/bce-embedding-base_v1')

	model.to(device)

	# get inputs
	inputs = tokenizer(text_list, padding=True, truncation=True, max_length=512, return_tensors="pt")
	inputs_on_device = {k: v.to(self.device) for k, v in inputs.items()}

	# get embeddings
	outputs = model(**inputs_on_device, return_dict=True)
	embeddings = outputs.last_hidden_state[:, 0] # cls pooler
	embeddings = embeddings / embeddings.norm(dim=1, keepdim=True) # normalize
	return embeddings
	def download_models():
	print("正在下载Luotuo-Bert")
	# Import our models. The package will take care of downloading the models automatically
	model_args = Namespace(do_mlm=None, pooler_type="cls", temp=0.05, mlp_only_train=False,
	init_embeddings_model=None)
	model = AutoModel.from_pretrained("silk-road/luotuo-bert-medium", trust_remote_code=True, model_args=model_args).to(
	device)
	print("Luotuo-Bert下载完毕")
	return model

	def get_luotuo_model():
	global _luotuo_model
	if _luotuo_model is None:
	_luotuo_model = download_models()
	return _luotuo_model


	def luotuo_embedding(model, texts):
	# Tokenize the texts_source
	tokenizer = AutoTokenizer.from_pretrained("silk-road/luotuo-bert-medium")
	inputs = tokenizer(texts, padding=True, truncation=False, return_tensors="pt")
	inputs = inputs.to(device)
	# Extract the embeddings
	# Get the embeddings
	with torch.no_grad():
	embeddings = model(**inputs, output_hidden_states=True, return_dict=True, sent_emb=True).pooler_output
	return embeddings

	def luotuo_en_embedding( texts ):
	# this function implemented by Cheng
	global _luotuo_model_en
	global _luotuo_en_tokenizer

	if _luotuo_model_en is None:
	_luotuo_en_tokenizer = AutoTokenizer.from_pretrained("silk-road/luotuo-bert-en")
	_luotuo_model_en = AutoModel.from_pretrained("silk-road/luotuo-bert-en").to(device)

	if _luotuo_en_tokenizer is None:
	_luotuo_en_tokenizer = AutoTokenizer.from_pretrained("silk-road/luotuo-bert-en")

	inputs = _luotuo_en_tokenizer(texts, padding=True, truncation=False, return_tensors="pt")
	inputs = inputs.to(device)

	with torch.no_grad():
	embeddings = _luotuo_model_en(**inputs, output_hidden_states=True, return_dict=True, sent_emb=True).pooler_output

	return embeddings


	def get_embedding_for_chinese(model, texts):
	model = model.to(device)
	# str or strList
	texts = texts if isinstance(texts, list) else [texts]
	# 截断
	for i in range(len(texts)):
	if len(texts[i]) > 510:
	texts[i] = texts[i][:510]
	if len(texts) >= 64:
	embeddings = []
	chunk_size = 64
	for i in range(0, len(texts), chunk_size):
	embeddings.append(luotuo_embedding(model, texts[i: i + chunk_size]))
	return torch.cat(embeddings, dim=0)
	else:
	return luotuo_embedding(model, texts)


	def is_chinese_or_english(text):
	# no longer use online openai api
	return "chinese"

	text = list(text)
	is_chinese, is_english = 0, 0

	for char in text:
	# 判断字符的Unicode值是否在中文字符的Unicode范围内
	if '\u4e00' <= char <= '\u9fa5':
	is_chinese += 4
	# 判断字符是否为英文字符（包括大小写字母和常见标点符号）
	elif ('\u0041' <= char <= '\u005a') or ('\u0061' <= char <= '\u007a'):
	is_english += 1
	if is_chinese >= is_english:
	return "chinese"
	else:
	return "english"


	def get_embedding_openai(text, model="text-embedding-ada-002"):
	text = text.replace("\n", " ")
	return client.embeddings.create(input = [text], model=model).data[0].embedding

	def get_embedding_for_english(text, model="text-embedding-ada-002"):
	text = text.replace("\n", " ")
	return client.embeddings.create(input = [text], model=model).data[0].embedding

	import os

	def foo_openai( text ):
	# dim = 1536

	openai_key = os.environ.get("OPENAI_API_KEY")

	if isinstance(texts, list):
	index = random.randint(0, len(texts) - 1)
	if openai_key is None or is_chinese_or_english(texts[index]) == "chinese":
	return [embed.cpu().tolist() for embed in get_embedding_for_chinese(get_luotuo_model(), texts)]
	else:
	return [get_embedding_for_english(text) for text in texts]
	else:
	if openai_key is None or is_chinese_or_english(texts) == "chinese":
	return get_embedding_for_chinese(get_luotuo_model(), texts)[0].cpu().tolist()
	else:
	return get_embedding_for_english(texts)


	### BGE family


	# ======== add bge_zh mmodel
	# by Cheng Li
	# 这一次我们试图一次性去适配更多的模型
	import torch

	_model_pool = {}
	_tokenizer_pool = {}

	# BAAI/bge-small-zh-v1.5

	def get_general_embeddings( sentences , model_name = "BAAI/bge-small-zh-v1.5" ):

	global _model_pool
	global _tokenizer_pool

	if model_name not in _model_pool:
	from transformers import AutoTokenizer, AutoModel
	_tokenizer_pool[model_name] = AutoTokenizer.from_pretrained(model_name)
	_model_pool[model_name] = AutoModel.from_pretrained(model_name).to(device)

	_model_pool[model_name].eval()

	# Tokenize sentences
	encoded_input = _tokenizer_pool[model_name](sentences, padding=True, truncation=True, return_tensors='pt', max_length = 512).to(device)

	# Compute token embeddings
	with torch.no_grad():
	model_output = _model_pool[model_name](**encoded_input)
	# Perform pooling. In this case, cls pooling.
	sentence_embeddings = model_output[0][:, 0]

	# normalize embeddings
	sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
	return sentence_embeddings.cpu().tolist()

	def get_general_embedding( text_or_texts , model_name = "BAAI/bge-small-zh-v1.5" ):
	if isinstance(text_or_texts, str):
	return get_general_embeddings([text_or_texts], model_name)[0]
	else:
	return get_general_embeddings_safe(text_or_texts, model_name)

	general_batch_size = 16

	import math

	def get_general_embeddings_safe(sentences, model_name = "BAAI/bge-small-zh-v1.5"):

	embeddings = []

	num_batches = math.ceil(len(sentences) / general_batch_size)

	from tqdm import tqdm

	for i in tqdm( range(num_batches) ):
	# print("run bge with batch ", i)
	start_index = i * general_batch_size
	end_index = min(len(sentences), start_index + general_batch_size)
	batch = sentences[start_index:end_index]
	embs = get_general_embeddings(batch, model_name)
	embeddings.extend(embs)

	return embeddings

	def get_bge_zh_embedding( text_or_texts ):
	return get_general_embedding(text_or_texts, "BAAI/bge-small-zh-v1.5")