refresh

get_dataset.py

@@ -1,68 +0,0 @@
-# coding=utf-8
-import sys
-sys.path.append("../")
-
-from collections import defaultdict
-from .utils import is_float, load_txt
-
-import random
-
-random.seed(1234)
-
-
-class CreateDataset:
-    def __init__(self, max_input_len=1500):
-        self.prompt = load_txt("../prompt/dataset_character.txt")
-        self.max_input_len = max_input_len  # 小于(seq-length)-(max-gen-length)
-        self.example_split_flag = f"\n{'-' * 20}\n"
-
-        self.dataset = defaultdict(list)
-        self.manual_dataset = []
-
-    @staticmethod
-    def choose_examples(similar_examples,
-                        max_length,
-                        train_flag=False,
-                        dialog=None,
-                        example_split_flag=f"\n{'-' * 20}\n"):
-        if isinstance(similar_examples, str):
-            new_similar_examples = [x.strip() for x in similar_examples.split(example_split_flag)]
-        else:
-            # 去重
-            new_similar_examples = []
-            for example in similar_examples:
-                if (isinstance(example, list) or isinstance(example, tuple)) and len(example) == 2 and is_float(
-                        example[0]):
-                    # 包含score
-                    example = example[1]
-
-                try:
-                    example = "\n".join(example).strip()
-                except TypeError:
-                    raise TypeError(f"example: {example}")
-                if train_flag and dialog and (example in dialog or dialog in example):
-                    continue
-
-                # example去重
-                if train_flag:
-                    # 部分相似也去掉
-                    flag = False
-                    for n_example in new_similar_examples:
-                        if example in n_example or n_example in example:
-                            flag = True
-                            break
-                    if not flag:
-                        new_similar_examples.append(example)
-                else:
-                    if example not in new_similar_examples:
-                        new_similar_examples.append(example)
-
-        results = []
-        total_length = 0
-        for example in new_similar_examples:
-            total_length += len(example) if not total_length else len(example_split_flag) + len(example)
-            if total_length > max_length:
-                break
-            results.append(example)
-        results = example_split_flag.join(results).strip()
-        return results

logger.py

@@ -1,60 +0,0 @@
-# coding=utf-8
-from logging.handlers import TimedRotatingFileHandler
-
-import os
-import sys
-import logging
-
-
-class LoggerFactory:
-
-    @staticmethod
-    def create_logger(name=None, level=logging.INFO):
-        """create a logger
-
-        Args:
-            name (str): name of the logger
-            level: level of logger
-
-        Raises:
-            ValueError is name is None
-        """
-
-        if name is None:
-            raise ValueError("name for logger cannot be None")
-
-        formatter = logging.Formatter("[%(asctime)s] [%(levelname)s] "
-                                      "[%(filename)s:%(lineno)d:%(funcName)s] %(message)s")
-
-        logger_ = logging.getLogger(name)
-        logger_.setLevel(level)
-        logger_.propagate = False
-        ch = logging.StreamHandler(stream=sys.stdout)
-        ch.setLevel(level)
-        ch.setFormatter(formatter)
-        logger_.addHandler(ch)
-        return logger_
-
-    @staticmethod
-    def create_logger_with_file(log_file_path: str = None, logger_level=logging.INFO):
-        logger_inner = logging.getLogger()
-        logger_inner.setLevel(logger_level)
-        logger_inner.propagate = True
-
-        formatter = logging.Formatter(fmt="[%(asctime)s] [%(filename)s:%(lineno)s - %(levelname)s] %(message)s",
-                                      datefmt="%Y-%m-%d %H:%M:%S")
-
-        # TimedRotatingFileHandler
-        if log_file_path:
-            basedir = os.path.dirname(log_file_path)
-            if not os.path.isdir(basedir):
-                os.makedirs(basedir, exist_ok=True)
-            handler_file = TimedRotatingFileHandler(log_file_path, when="d", interval=1, backupCount=30)
-            handler_file.setFormatter(formatter)
-            logger_inner.addHandler(handler_file)
-
-        # StreamHandler
-        handler_console = logging.StreamHandler()
-        handler_console.setFormatter(formatter)
-        logger_inner.addHandler(handler_console)
-        return logger_inner

prompt_concat.py

@@ -1,170 +0,0 @@
-# coding=utf-8
-from copy import deepcopy
-from .get_dataset import CreateDataset
-from .logger import LoggerFactory
-from .retrieve_dialog import RetrieveDialog
-from .utils import load_json, load_txt, save_to_json
-
-import logging
-import os
-
-logger = LoggerFactory.create_logger(name="test", level=logging.INFO)
-
-
-class GetManualTestSamples:
-    def __init__(
-        self,
-        role_name,
-        role_data_path,
-        save_samples_dir,
-        save_samples_path=None,
-        prompt_path="dataset_character.txt",
-        max_seq_len=4000,
-        retrieve_num=20,
-    ):
-        self.role_name = role_name.strip()
-        self.role_data = load_json(role_data_path)
-        self.role_info = self.role_data[0]["role_info"].strip()
-
-        self.prompt = load_txt(prompt_path)
-        self.prompt = self.prompt.replace("${role_name}", self.role_name)
-        self.prompt = self.prompt.replace("${role_info}",
-                                          f"以下是{self.role_name}的人设：\n{self.role_info}\n").strip()
-
-        self.retrieve_num = retrieve_num
-        self.retrieve = RetrieveDialog(role_name=self.role_name,
-                                       raw_dialog_list=[d["dialog"] for d in self.role_data],
-                                       retrieve_num=retrieve_num)
-
-        self.max_seq_len = max_seq_len
-        if not save_samples_path:
-            save_samples_path = f"{self.role_name}.json"
-        self.save_samples_path = os.path.join(save_samples_dir, save_samples_path)
-
-    def _add_simi_dialog(self, history: list, content_length):
-        retrieve_results = self.retrieve.get_retrieve_res(history, self.retrieve_num)
-        simi_dialogs = deepcopy(retrieve_results)
-
-        if simi_dialogs:
-            simi_dialogs = CreateDataset.choose_examples(simi_dialogs,
-                                                         max_length=self.max_seq_len - content_length,
-                                                         train_flag=False)
-        logger.debug(f"retrieve_results: {retrieve_results}\nsimi_dialogs: {simi_dialogs}.")
-        return simi_dialogs, retrieve_results
-
-    def get_qa_samples_by_file(self,
-                               questions_path,
-                               user_name="user",
-                               keep_retrieve_results_flag=False
-                               ):
-        questions = load_txt(questions_path).splitlines()
-        samples = []
-        for question in questions:
-            question = question.replace('\\n', "\n")
-            query = f"{user_name}:{question}" if ":" not in question else question
-            content = self.prompt.replace("${dialog}", query)
-            content = content.replace("${user_name}", user_name).strip()
-
-            history = [query]
-            simi_dialogs, retrieve_results = self._add_simi_dialog(history, len(content))
-
-            sample = {
-                "role_name": self.role_name,
-                "role_info": self.role_info,
-                "user_name": user_name,
-                "dialog": history,
-                "simi_dialogs": simi_dialogs,
-            }
-            if keep_retrieve_results_flag and retrieve_results:
-                sample["retrieve_results"] = retrieve_results
-            samples.append(sample)
-        self._save_samples(samples)
-
-    def get_qa_samples_by_query(self,
-                                questions_query,
-                                user_name="user",
-                                keep_retrieve_results_flag=False
-                                ):
-        question = questions_query
-        samples = []
-        question = question.replace('\\n', "\n")
-        query = f"{user_name}: {question}" if ":" not in question else question
-        content = self.prompt.replace("${dialog}", query)
-        content = content.replace("${user_name}", user_name).strip()
-
-        history = [query]
-        simi_dialogs, retrieve_results = self._add_simi_dialog(history, len(content))
-
-        sample = {
-            "role_name": self.role_name,
-            "role_info": self.role_info,
-            "user_name": user_name,
-            "dialog": history,
-            "simi_dialogs": simi_dialogs,
-        }
-        if keep_retrieve_results_flag and retrieve_results:
-            sample["retrieve_results"] = retrieve_results
-        samples.append(sample)
-        self._save_samples(samples)
-
-    def _save_samples(self, samples):
-        data = samples
-        save_to_json(data, self.save_samples_path)
-
-
-class CreateTestDataset:
-    def __init__(self,
-                 role_name,
-                 role_samples_path=None,
-                 role_data_path=None,
-                 prompt_path="dataset_character.txt",
-                 max_seq_len=4000):
-        self.max_seq_len = max_seq_len
-        self.role_name = role_name
-
-        self.prompt = load_txt(prompt_path)
-        self.prompt = self.prompt.replace("${role_name}", role_name).strip()
-
-        if not role_data_path:
-            print("need role_data_path, check please!")
-        self.default_simi_dialogs = None
-        if os.path.exists(role_data_path):
-            data = load_json(role_data_path)
-            role_info = data[0]["role_info"]
-        else:
-            raise ValueError(f"{self.role_name} didn't find role_info.")
-        self.role_info = role_info
-        self.prompt = self.prompt.replace("${role_info}", f"以下是{self.role_name}的人设：\n{self.role_info}\n").strip()
-
-        if role_samples_path:
-            self.role_samples_path = role_samples_path
-        else:
-            print("check role_samples_path please!")
-
-    def load_samples(self):
-        samples = load_json(self.role_samples_path)
-        results = []
-        for sample in samples:
-            input_text = self.prompt
-
-            simi_dialogs = sample.get("simi_dialogs", None)
-            if not simi_dialogs:
-                simi_dialogs = self.default_simi_dialogs
-            if not simi_dialogs:
-                raise ValueError(f"didn't find simi_dialogs.")
-            simi_dialogs = CreateDataset.choose_examples(simi_dialogs,
-                                                         max_length=self.max_seq_len - len(input_text),
-                                                         train_flag=False)
-
-            input_text = input_text.replace("${simi_dialog}", simi_dialogs)
-            user_name = sample.get("user_name", "user")
-            input_text = input_text.replace("${user_name}", user_name)
-
-            dialog = "\n".join(sample["dialog"]) if isinstance(sample["dialog"], list) else sample["dialog"]
-            input_text = input_text.replace("${dialog}", dialog)
-
-            assert len(input_text) < self.max_seq_len
-            results.append({
-                "input_text": input_text,
-            })
-        return results

retrieve_dialog.py

@@ -1,135 +0,0 @@
-# coding=utf-8
-from sentence_transformers import SentenceTransformer
-from .utils import load_json
-
-import faiss
-import logging
-import os
-import re
-import torch
-
-logger = logging.getLogger(__name__)
-
-
-class RetrieveDialog:
-    def __init__(self,
-                 role_name,
-                 raw_dialog_list: list = None,
-                 retrieve_num=20,
-                 min_mean_role_utter_length=10):
-        if torch.cuda.is_available():
-            gpu_id = 0
-            torch.cuda.set_device(gpu_id)
-
-        assert raw_dialog_list
-
-        self.role_name = role_name
-        self.min_mean_role_utter_length = min_mean_role_utter_length
-        self.retrieve_num = retrieve_num
-
-        # config = load_json("config/config.json")
-        # local_dir = config["bge_local_path"]
-        local_dir = os.environ.get('MODEL_PATH', 'IndexTeam/Index-1.9B-Character')
-
-        if not os.path.exists(local_dir):
-            print("Please download bge-large-zh-v1.5 first!")
-        self.emb_model = SentenceTransformer(local_dir)
-
-        self.dialogs, self.context_index = self._get_emb_base_by_list(raw_dialog_list)
-
-        logger.info(f"dialog db num: {len(self.dialogs)}")
-        logger.info(f"RetrieveDialog init success.")
-
-    @staticmethod
-    def dialog_preprocess(dialog: list, role_name):
-        dialog_new = []
-        # 把人名替换掉，减少对检索的影响
-        user_names = []
-        role_utter_length = []
-        for num in range(len(dialog)):
-            utter = dialog[num]
-            try:
-                user_name, utter_txt = re.split('[:：]', utter, maxsplit=1)
-            except ValueError as e:
-                logging.error(f"utter:{utter} can't find user_name.")
-                return None, None
-
-            if user_name != role_name:
-                if user_name not in user_names:
-                    user_names.append(user_name)
-                index = user_names.index(user_name)
-                utter = utter.replace(user_name, f"user{index}", 1)
-            else:
-                role_utter_length.append(len(utter_txt))
-            dialog_new.append(utter)
-        return dialog_new, user_names, role_utter_length
-
-    def _get_emb_base_by_list(self, raw_dialog_list):
-        logger.info(f"raw dialog db num: {len(raw_dialog_list)}")
-        new_raw_dialog_list = []
-        context_list = []
-
-        # 为了兼容因为句长把所有对话都过滤掉的情况
-        new_raw_dialog_list_total = []
-        context_list_total = []
-        for raw_dialog in raw_dialog_list:
-            if not raw_dialog:
-                continue
-
-            end = 0
-            for x in raw_dialog[::-1]:
-                if x.startswith(self.role_name):
-                    break
-                end += 1
-
-            raw_dialog = raw_dialog[:len(raw_dialog) - end]
-            new_dialog, user_names, role_utter_length = self.dialog_preprocess(raw_dialog, self.role_name)
-            if not new_dialog or not role_utter_length:
-                continue
-
-            if raw_dialog in new_raw_dialog_list_total:
-                continue
-
-            # 获得embedding时，不需要最后一句答案
-            context = "\n".join(new_dialog) if len(new_dialog) < 2 else "\n".join(new_dialog[:-1])
-
-            new_raw_dialog_list_total.append(raw_dialog)
-            context_list_total.append(context)
-
-            # 句长过滤
-            role_length_mean = sum(role_utter_length) / len(role_utter_length)
-            if role_length_mean < self.min_mean_role_utter_length:
-                continue
-            new_raw_dialog_list.append(raw_dialog)
-            context_list.append(context)
-
-        assert len(new_raw_dialog_list) == len(context_list)
-        logger.debug(f"new_raw_dialog num: {len(new_raw_dialog_list)}")
-
-        # 兼容样本过少的情况
-        if len(new_raw_dialog_list) < self.retrieve_num:
-            new_raw_dialog_list = new_raw_dialog_list_total
-            context_list = context_list_total
-
-        # 对话向量库
-        context_vectors = self.emb_model.encode(context_list, normalize_embeddings=True)
-        context_index = faiss.IndexFlatL2(context_vectors.shape[1])
-        context_index.add(context_vectors)
-
-        return new_raw_dialog_list, context_index
-
-    def get_retrieve_res(self, dialog: list, retrieve_num: int):
-        logger.debug(f"dialog: {dialog}")
-
-        # 同样去掉user name影响
-        dialog, _, _ = self.dialog_preprocess(dialog, self.role_name)
-        dialog_vector = self.emb_model.encode(["\n".join(dialog)], normalize_embeddings=True)
-
-        simi_dialog_distance, simi_dialog_index = self.context_index.search(
-            dialog_vector, min(retrieve_num, len(self.dialogs)))
-        simi_dialog_results = [
-            (str(simi_dialog_distance[0][num]), self.dialogs[index]) for num, index in enumerate(simi_dialog_index[0])
-        ]
-        logger.debug(f"dialog retrieve res: {simi_dialog_results}")
-
-        return simi_dialog_results

src/retrieve_dialog.py

@@ -27,8 +27,9 @@ class RetrieveDialog:
         self.min_mean_role_utter_length = min_mean_role_utter_length
         self.retrieve_num = retrieve_num
 
-        config = load_json("config/config.json")
-        local_dir = config["bge_local_path"]
+        # config = load_json("config/config.json")
+        # local_dir = config["bge_local_path"]
+        local_dir = os.environ.get('MODEL_PATH', 'IndexTeam/Index-1.9B-Character')
 
         if not os.path.exists(local_dir):
             print("Please download bge-large-zh-v1.5 first!")

utils.py

@@ -1,59 +0,0 @@
-# coding=utf-8
-import csv
-import json
-import os
-
-
-def read_csv_to_json(file_path, role_name, role_info):
-    json_list = []
-    
-    with open(file_path, mode="r", newline="", encoding="utf-8") as csvfile:
-        csv_reader = csv.reader(csvfile)
-        _ = next(csv_reader)
-        
-        for row in csv_reader:
-            json_object = {
-                "role_name": role_name,
-                "role_info": role_info,
-                "dialog": row[1].split("\n"),
-            }
-            json_list.append(json_object)
-    
-    return json_list
-
-
-def save_json(json_list, output_path):
-    with open(output_path, "w", encoding="utf-8") as jsonfile:
-        json.dump(json_list, jsonfile, ensure_ascii=False, indent=4)
-
-
-def decode_csv_to_json(role_data_path, role_name, role_info, json_output_path):
-    json_data = read_csv_to_json(role_data_path, role_name, role_info)
-    save_json(json_data, json_output_path)
-    
-
-def load_txt(path):
-    with open(path, "r", encoding="utf-8", errors="ignore") as file:
-        text = file.read()
-    return text
-
-
-def load_json(path):
-    with open(path, "r", encoding="utf-8") as f:
-        data = json.load(f)
-    return data
-
-
-def save_to_json(data, filepath, flag="w"):
-    if not os.path.exists(os.path.dirname(filepath)):
-        os.makedirs(os.path.dirname(filepath))
-    with open(filepath, flag, encoding="utf-8") as f:
-        f.write(json.dumps(data, ensure_ascii=False, indent=3))
-
-
-def is_float(my_str):
-    try:
-        num = float(my_str)
-        return True
-    except ValueError:
-        return False