Upload 6 files

__init__.py

@@ -0,0 +1,6 @@
+# coding=utf-8
+from .get_dataset import *
+from .logger import *
+from .prompt_concat import *
+from .retrieve_dialog import *
+from .utils import *

get_dataset.py

@@ -0,0 +1,68 @@
+# 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

@@ -0,0 +1,60 @@
+# 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

@@ -0,0 +1,170 @@
+# 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

utils.py

@@ -0,0 +1,59 @@
+# 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