from .utils import base64_to_float_array, base64_to_string def get_text_from_data( data ): if "text" in data: return data['text'] elif "enc_text" in data: # from .utils import base64_to_string return base64_to_string( data['enc_text'] ) else: print("warning! failed to get text from data ", data) return "" def parse_rag(text): lines = text.split("\n") ans = [] for i, line in enumerate(lines): if "{{RAG对话}}" in line: ans.append({"n": 1, "max_token": -1, "query": "default", "lid": i}) elif "{{RAG对话|" in line: query_info = line.split("|")[1].rstrip("}}") ans.append({"n": 1, "max_token": -1, "query": query_info, "lid": i}) elif "{{RAG多对话|" in line: parts = line.split("|") max_token = int(parts[1].split("<=")[1]) max_n = int(parts[2].split("<=")[1].rstrip("}}")) ans.append({"n": max_n, "max_token": max_token, "query": "default", "lid": i}) return ans class ChatHaruhi: def __init__(self, role_name = None, user_name = None, persona = None, stories = None, story_vecs = None, role_from_hf = None, role_from_jsonl = None, llm = None, # 默认的message2response的函数 llm_async = None, # 默认的message2response的async函数 user_name_in_message = "default", verbose = None, embed_name = None, embedding = None, db = None, token_counter = "default", max_input_token = 1800, max_len_story_haruhi = 1000, max_story_n_haruhi = 5 ): self.verbose = True if verbose is None or verbose else False self.db = db self.embed_name = embed_name self.max_len_story_haruhi = max_len_story_haruhi # 这个设置只对过往Haruhi的sugar角色有效 self.max_story_n_haruhi = max_story_n_haruhi # 这个设置只对过往Haruhi的sugar角色有效 self.last_query_msg = None if embedding is None: self.embedding = self.set_embedding_with_name( embed_name ) if persona and role_name and stories and story_vecs and len(stories) == len(story_vecs): # 完全从外部设置，这个时候要求story_vecs和embedding的返回长度一致 self.persona, self.role_name, self.user_name = persona, role_name, user_name self.build_db(stories, story_vecs) elif persona and role_name and stories: # 从stories中提取story_vecs，重新用self.embedding进行embedding story_vecs = self.extract_story_vecs(stories) self.persona, self.role_name, self.user_name = persona, role_name, user_name self.build_db(stories, story_vecs) elif role_from_hf: # 从hf加载role self.persona, new_role_name, self.stories, self.story_vecs = self.load_role_from_hf(role_from_hf) if new_role_name: self.role_name = new_role_name else: self.role_name = role_name self.user_name = user_name self.build_db(self.stories, self.story_vecs) elif role_from_jsonl: # 从jsonl加载role self.persona, new_role_name, self.stories, self.story_vecs = self.load_role_from_jsonl(role_from_jsonl) if new_role_name: self.role_name = new_role_name else: self.role_name = role_name self.user_name = user_name self.build_db(self.stories, self.story_vecs) elif persona and role_name: # 这个时候也就是说没有任何的RAG， self.persona, self.role_name, self.user_name = persona, role_name, user_name self.db = None elif role_name and self.check_sugar( role_name ): # 这个时候是sugar的role self.persona, self.role_name, self.stories, self.story_vecs = self.load_role_from_sugar( role_name ) self.build_db(self.stories, self.story_vecs) # 与 江YH讨论 所有的载入方式都要在外部使用 add_rag_prompt_after_persona() 防止混淆 # self.add_rag_prompt_after_persona() else: raise ValueError("persona和role_name必须同时设置，或者role_name是ChatHaruhi的预设人物") self.llm, self.llm_async = llm, llm_async if not self.llm and self.verbose: print("warning, llm没有设置，仅get_message起作用，调用chat将回复idle message") self.user_name_in_message = user_name_in_message self.previous_user_pool = set([user_name]) if user_name else set() self.current_user_name_in_message = user_name_in_message.lower() == "add" self.idle_message = "idel message, you see this because self.llm has not been set." if token_counter.lower() == "default": # TODO change load from util from .utils import tiktoken_counter self.token_counter = tiktoken_counter elif token_counter == None: self.token_counter = lambda x: 0 else: self.token_counter = token_counter if self.verbose: print("user set costomized token_counter") self.max_input_token = max_input_token self.history = [] def check_sugar(self, role_name): from .sugar_map import sugar_role_names, enname2zhname return role_name in sugar_role_names def load_role_from_sugar(self, role_name): from .sugar_map import sugar_role_names, enname2zhname en_role_name = sugar_role_names[role_name] new_role_name = enname2zhname[en_role_name] role_from_hf = "silk-road/ChatHaruhi-RolePlaying/" + en_role_name persona, _, stories, story_vecs = self.load_role_from_hf(role_from_hf) return persona, new_role_name, stories, story_vecs def add_rag_prompt_after_persona( self ): rag_sentence = "{{RAG多对话|token<=" + str(self.max_len_story_haruhi) + "|n<=" + str(self.max_story_n_haruhi) + "}}" self.persona += "Classic scenes for the role are as follows:\n" + rag_sentence + "\n" def set_embedding_with_name(self, embed_name): if embed_name is None or embed_name == "bge_zh": from .embeddings import get_bge_zh_embedding self.embed_name = "bge_zh" return get_bge_zh_embedding elif embed_name == "foo": from .embeddings import foo_embedding return foo_embedding elif embed_name == "bce": from .embeddings import foo_bce return foo_bce elif embed_name == "openai" or embed_name == "luotuo_openai": from .embeddings import foo_openai return foo_openai def set_new_user(self, user): if len(self.previous_user_pool) > 0 and user not in self.previous_user_pool: if self.user_name_in_message.lower() == "default": if self.verbose: print(f'new user {user} included in conversation') self.current_user_name_in_message = True self.user_name = user self.previous_user_pool.add(user) def chat(self, user, text): self.set_new_user(user) message = self.get_message(user, text) if self.llm: response = self.llm(message) self.append_message(response) return response return None async def async_chat(self, user, text): self.set_new_user(user) message = self.get_message(user, text) if self.llm_async: response = await self.llm_async(message) self.append_message(response) return response def parse_rag_from_persona(self, persona, text = None): #每个query_rag需要饱含 # "n" 需要几个story # "max_token" 最多允许多少个token，如果-1则不限制 # "query" 需要查询的内容，如果等同于"default"则替换为text # "lid" 需要替换的行，这里直接进行行替换，忽视行的其他内容 query_rags = parse_rag( persona ) if text is not None: for rag in query_rags: if rag['query'] == "default": rag['query'] = text return query_rags, self.token_counter(persona) def append_message( self, response , speaker = None ): if self.last_query_msg is not None: self.history.append(self.last_query_msg) self.last_query_msg = None if speaker is None: # 如果role是none，则认为是本角色{{role}}输出的句子 self.history.append({"speaker":"{{role}}","content":response}) # 叫speaker是为了和role进行区分 else: self.history.append({"speaker":speaker,"content":response}) def check_recompute_stories_token(self): return len(self.db.metas) == len(self.db.stories) def recompute_stories_token(self): self.db.metas = [self.token_counter(story) for story in self.db.stories] def rag_retrieve( self, query, n, max_token, avoid_ids = [] ): # 返回一个rag_id的列表 query_vec = self.embedding(query) self.db.clean_flag() self.db.disable_story_with_ids( avoid_ids ) retrieved_ids = self.db.search( query_vec, n ) if self.check_recompute_stories_token(): self.recompute_stories_token() sum_token = 0 ans = [] for i in range(0, len(retrieved_ids)): if i == 0: sum_token += self.db.metas[retrieved_ids[i]] ans.append(retrieved_ids[i]) continue else: sum_token += self.db.metas[retrieved_ids[i]] if sum_token <= max_token: ans.append(retrieved_ids[i]) else: break return ans def rag_retrieve_all( self, query_rags, rest_limit ): # 返回一个rag_ids的列表 retrieved_ids = [] rag_ids = [] for query_rag in query_rags: query = query_rag['query'] n = query_rag['n'] max_token = rest_limit if rest_limit > query_rag['max_token'] and query_rag['max_token'] > 0: max_token = query_rag['max_token'] rag_id = self.rag_retrieve( query, n, max_token, avoid_ids = retrieved_ids ) rag_ids.append( rag_id ) retrieved_ids += rag_id return rag_ids def append_history_under_limit(self, message, rest_limit): # 返回一个messages的列表 # print("call append history_under_limit") # 从后往前计算token，不超过rest limit， # 如果speaker是{{role}J,则message的role是assistant current_limit = rest_limit history_list = [] for item in reversed(self.history): current_token = self.token_counter(item['content']) current_limit -= current_token if current_limit < 0: break else: history_list.append(item) history_list = list(reversed(history_list)) # TODO: 之后为了解决多人对话，这了content还会额外增加speaker: content这样的信息 for item in history_list: if item['speaker'] == "{{role}}": message.append({"role":"assistant","content":item['content']}) else: message.append({"role":"user","content":item['content']}) return message def get_message(self, user, text): query_token = self.token_counter(text) # 首先获取需要多少个rag story query_rags, persona_token = self.parse_rag_from_persona( self.persona, text ) #每个query_rag需要饱含 # "n" 需要几个story # "max_token" 最多允许多少个token，如果-1则不限制 # "query" 需要查询的内容，如果等同于"default"则替换为text # "lid" 需要替换的行，这里直接进行行替换，忽视行的其他内容 rest_limit = self.max_input_token - persona_token - query_token if self.verbose: print(f"query_rags: {query_rags} rest_limit = { rest_limit }") rag_ids = self.rag_retrieve_all( query_rags, rest_limit ) # 将rag_ids对应的故事 替换到persona中 augmented_persona = self.augment_persona( self.persona, rag_ids, query_rags ) system_prompt = self.package_system_prompt( self.role_name, augmented_persona ) token_for_system = self.token_counter( system_prompt ) rest_limit = self.max_input_token - token_for_system - query_token message = [{"role":"system","content":system_prompt}] message = self.append_history_under_limit( message, rest_limit ) # TODO: 之后为了解决多人对话，这了content还会额外增加speaker: content这样的信息 message.append({"role":"user","content":text}) self.last_query_msg = {"speaker":user,"content":text} return message def package_system_prompt(self, role_name, augmented_persona): bot_name = role_name return f"""You are now in roleplay conversation mode. Pretend to be {bot_name} whose persona follows: {augmented_persona} You will stay in-character whenever possible, and generate responses as if you were {bot_name}""" def augment_persona(self, persona, rag_ids, query_rags): lines = persona.split("\n") for rag_id, query_rag in zip(rag_ids, query_rags): lid = query_rag['lid'] new_text = "" for id in rag_id: new_text += "###\n" + self.db.stories[id].strip() + "\n" new_text = new_text.strip() lines[lid] = new_text return "\n".join(lines) def load_role_from_jsonl( self, role_from_jsonl ): import json datas = [] with open(role_from_jsonl, 'r') as f: for line in f: try: datas.append(json.loads(line)) except: continue column_name = "" from .embeddings import embedname2columnname if self.embed_name in embedname2columnname: column_name = embedname2columnname[self.embed_name] else: print('warning! unkown embedding name ', self.embed_name ,' while loading role') column_name = 'luotuo_openai' stories, story_vecs, persona = self.extract_text_vec_from_datas(datas, column_name) return persona, None, stories, story_vecs def load_role_from_hf(self, role_from_hf): # 从hf加载role # self.persona, new_role_name, self.stories, self.story_vecs = self.load_role_from_hf(role_from_hf) from datasets import load_dataset if role_from_hf.count("/") == 1: dataset = load_dataset(role_from_hf) datas = dataset["train"] elif role_from_hf.count("/") >= 2: split_index = role_from_hf.index('/') second_split_index = role_from_hf.index('/', split_index+1) dataset_name = role_from_hf[:second_split_index] split_name = role_from_hf[second_split_index+1:] fname = split_name + '.jsonl' dataset = load_dataset(dataset_name,data_files={'train':fname}) datas = dataset["train"] column_name = "" from .embeddings import embedname2columnname if self.embed_name in embedname2columnname: column_name = embedname2columnname[self.embed_name] else: print('warning! unkown embedding name ', self.embed_name ,' while loading role') column_name = 'luotuo_openai' stories, story_vecs, persona = self.extract_text_vec_from_datas(datas, column_name) return persona, None, stories, story_vecs def extract_text_vec_from_datas(self, datas, column_name): # 从datas中提取text和vec # extract text and vec from huggingface dataset # return texts, vecs # from .utils import base64_to_float_array texts = [] vecs = [] for data in datas: if data[column_name] == 'system_prompt': system_prompt = get_text_from_data( data ) elif data[column_name] == 'config': pass else: vec = base64_to_float_array( data[column_name] ) text = get_text_from_data( data ) vecs.append( vec ) texts.append( text ) return texts, vecs, system_prompt def extract_story_vecs(self, stories): # 从stories中提取story_vecs if self.verbose: print(f"re-extract vector for {len(stories)} stories") story_vecs = [] from .embeddings import embedshortname2model_name from .embeddings import device if device.type != "cpu" and self.embed_name in embedshortname2model_name: # model_name = "BAAI/bge-small-zh-v1.5" model_name = embedshortname2model_name[self.embed_name] from .utils import get_general_embeddings_safe story_vecs = get_general_embeddings_safe( stories, model_name = model_name ) # 使用batch的方式进行embedding，非常快 else: from tqdm import tqdm for story in tqdm(stories): story_vecs.append(self.embedding(story)) return story_vecs def build_db(self, stories, story_vecs): # db的构造函数 if self.db is None: from .NaiveDB import NaiveDB self.db = NaiveDB() self.db.build_db(stories, story_vecs)