Spaces:
Running
Running
| import collections | |
| import os | |
| import pickle | |
| from argparse import Namespace | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from torch import cosine_similarity | |
| from transformers import AutoTokenizer, AutoModel | |
| def download_models(): | |
| # 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", trust_remote_code=True, model_args=model_args) | |
| return model | |
| class Text: | |
| def __init__(self, text_dir, model, num_steps, text_image_pkl_path=None, dict_text_pkl_path=None, pkl_path=None, dict_path=None, image_path=None, maps_path=None): | |
| self.dict_text_pkl_path = dict_text_pkl_path | |
| self.text_image_pkl_path = text_image_pkl_path | |
| self.text_dir = text_dir | |
| self.model = model | |
| self.num_steps = num_steps | |
| self.pkl_path = pkl_path | |
| self.dict_path = dict_path | |
| self.image_path = image_path | |
| self.maps_path = maps_path | |
| def get_embedding(self, texts): | |
| tokenizer = AutoTokenizer.from_pretrained("silk-road/luotuo-bert") | |
| model = download_models() | |
| # 截断 | |
| # str or strList | |
| texts = texts if isinstance(texts, list) else [texts] | |
| for i in range(len(texts)): | |
| if len(texts[i]) > self.num_steps: | |
| texts[i] = texts[i][:self.num_steps] | |
| # Tokenize the texts | |
| inputs = tokenizer(texts, padding=True, truncation=False, return_tensors="pt") | |
| # 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[0] if len(texts) == 1 else embeddings | |
| def read_text(self, save_embeddings=False, save_maps=False): | |
| """抽取、预存""" | |
| text_embeddings = collections.defaultdict() | |
| text_keys = [] | |
| dirs = os.listdir(self.text_dir) | |
| data = [] | |
| texts = [] | |
| id = 0 | |
| for dir in dirs: | |
| with open(self.text_dir + '/' + dir, 'r') as fr: | |
| for line in fr.readlines(): | |
| category = collections.defaultdict(str) | |
| ch = '：' if '：' in line else ':' | |
| if '旁白' in line: | |
| text = line.strip().split(ch)[1].strip() | |
| else: | |
| text = ''.join(list(line.strip().split(ch)[1])[1:-1]) # 提取「」内的文本 | |
| if text in text_keys: # 避免重复的text，导致embeds 和 maps形状不一致 | |
| continue | |
| text_keys.append(text) | |
| if save_maps: | |
| category["titles"] = dir.split('.')[0] | |
| category["id"] = str(id) | |
| category["text"] = text | |
| id = id + 1 | |
| data.append(dict(category)) | |
| texts.append(text) | |
| embeddings = self.get_embedding(texts) | |
| if save_embeddings: | |
| for text, embed in zip(texts, embeddings): | |
| text_embeddings[text] = self.get_embedding(text) | |
| if save_embeddings: | |
| self.store(self.pkl_path, text_embeddings) | |
| if save_maps: | |
| self.store(self.maps_path, data) | |
| return text_embeddings, data | |
| def load(self, load_pkl=False, load_maps=False, load_dict_text=False, load_text_image=False): | |
| if self.pkl_path and load_pkl: | |
| with open(self.pkl_path, 'rb') as f: | |
| return pickle.load(f) | |
| elif self.maps_path and load_maps: | |
| with open(self.maps_path, 'rb') as f: | |
| return pickle.load(f) | |
| elif self.dict_text_pkl_path and load_dict_text: | |
| with open(self.dict_text_pkl_path, 'rb') as f: | |
| return pickle.load(f) | |
| elif self.text_image_pkl_path and load_text_image: | |
| with open(self.text_image_pkl_path, 'rb') as f: | |
| return pickle.load(f) | |
| else: | |
| print("No pkl_path") | |
| def get_cosine_similarity(self, texts, get_image=False, get_texts=False): | |
| """ | |
| 计算文本列表的相似度避免重复计算query_similarity | |
| texts[0] = query | |
| """ | |
| if get_image: | |
| pkl = self.load(load_dict_text=True) | |
| elif get_texts: | |
| pkl = self.load(load_pkl=True) | |
| else: | |
| pkl = {} | |
| embeddings = self.get_embedding(texts[1:]).reshape(-1, 1536) | |
| for text, embed in zip(texts, embeddings): | |
| pkl[text] = embed | |
| query_embedding = self.get_embedding(texts[0]).reshape(1, -1) | |
| texts_embeddings = np.array([value.numpy().reshape(-1, 1536) for value in pkl.values()]).squeeze(1) | |
| return cosine_similarity(query_embedding, torch.from_numpy(texts_embeddings)) | |
| def store(self, path, data): | |
| with open(path, 'wb+') as f: | |
| pickle.dump(data, f) | |
| def text_to_image(self, text, save_dict_text=False): | |
| """ | |
| 给定文本出图片 | |
| 计算query 和 texts 的相似度，取最高的作为new_query 查询image | |
| 到text_image_dict 读取图片名 | |
| 然后到images里面加载该图片然后返回 | |
| """ | |
| if save_dict_text: | |
| text_image = {} | |
| with open(self.dict_path, 'r') as f: | |
| data = f.readlines() | |
| for sub_text, image in zip(data[::2], data[1::2]): | |
| text_image[sub_text.strip()] = image.strip() | |
| self.store(self.text_image_pkl_path, text_image) | |
| keys_embeddings = {} | |
| embeddings = self.get_embedding(list(text_image.keys())) | |
| for key, embed in zip(text_image.keys(), embeddings): | |
| keys_embeddings[key] = embed | |
| self.store(self.dict_text_pkl_path, keys_embeddings) | |
| if self.dict_path and self.image_path: | |
| # 加载 text-imageName | |
| text_image = self.load(load_text_image=True) | |
| keys = list(text_image.keys()) | |
| keys.insert(0, text) | |
| query_similarity = self.get_cosine_similarity(keys, get_image=True) | |
| key_index = query_similarity.argmax(dim=0) | |
| text = list(text_image.keys())[key_index] | |
| image = text_image[text] + '.jpg' | |
| if image in os.listdir(self.image_path): | |
| res = Image.open(self.image_path + '/' + image) | |
| # res.show() | |
| return res | |
| else: | |
| print("Image doesn't exist") | |
| else: | |
| print("No path") | |
| def text_to_text(self, text): | |
| pkl = self.load(load_pkl=True) | |
| texts = list(pkl.keys()) | |
| texts.insert(0, text) | |
| texts_similarity = self.get_cosine_similarity(texts, get_texts=True) | |
| key_index = texts_similarity.argmax(dim=0).item() | |
| value = list(pkl.keys())[key_index] | |
| return value | |
| # if __name__ == '__main__': | |
| # pkl_path = './pkl/texts.pkl' | |
| # maps_path = './pkl/maps.pkl' | |
| # text_image_pkl_path='./pkl/text_image.pkl' | |
| # dict_path = "../characters/haruhi/text_image_dict.txt" | |
| # dict_text_pkl_path = './pkl/dict_text.pkl' | |
| # image_path = "../characters/haruhi/images" | |
| # text_dir = "../characters/haruhi/texts" | |
| # model = download_models() | |
| # text = Text(text_dir, text_image_pkl_path=text_image_pkl_path, maps_path=maps_path, | |
| # dict_text_pkl_path=dict_text_pkl_path, model=model, num_steps=50, pkl_path=pkl_path, | |
| # dict_path=dict_path, image_path=image_path) | |
| # text.read_text(save_maps=True, save_embeddings=True) | |
| # data = text.load(load_pkl=True) | |
| # sub_text = "你好！" | |
| # image = text.text_to_image(sub_text) | |
| # print(image) | |
| # sub_texts = ["hello", "你好"] | |
| # print(text.get_cosine_similarity(sub_texts)) | |
| # value = text.text_to_text(sub_text) | |
| # print(value) | |