update

Author: neavo

00.py

@@ -32,15 +32,15 @@
 from model.NERTrainerCallback import NERTrainerCallback
 
 # 参数设置
-MODEL_NAME = "facebookai_xlm_roberta_base_pretrain_20240823"
+MODEL_NAME = "facebookai_xlm_roberta_base_pretrain_20240826"
 MODEL_PATH = f"assets/{MODEL_NAME}"
 OUTPUT_PATH = "output"
 DATASET_PATH = "dataset/ner"
 EPOCHS = 24
 PATIENCE = 12
 PATIENCE_KEEPER = 0
 BATCH_SIZE = 32
-GRADIENT_ACCUMULATION_SIZE = 64
+GRADIENT_ACCUMULATION_SIZE = 32
 FROZEN_LAYER = 0
 LEARNING_RATE = 2 * 1e-5
 DO_LOWER_CASE = False
@@ -65,13 +65,13 @@ def load_dataset(tokenizer):
         if file.name.endswith(".json"):
             with open(file.path, "r", encoding = "utf-8") as file:
                 count = count + 1
-                datas.extend(json.load(file))
+                datas.extend(random.sample(json.load(file), 10000))
 
     print(f"")
     print(f"找到数据文件 {count} 个，共 {len(datas)} 条数据 ...")
 
     # 分割数据集
-    train_datas, test_datas = train_test_split(datas, test_size = 0.02, shuffle = True, random_state = 42)
+    train_datas, test_datas = train_test_split(datas, test_size = 0.025, shuffle = True, random_state = 42)
 
     # 创建数据集和数据加载器
     print(f"")

01.py

@@ -41,9 +41,10 @@
     ("dataset/pretrain/en_r18_visual_novels", 20 * 10000),
     ("dataset/pretrain/zh", 20 * 10000),
     ("dataset/pretrain/zh_r18_pixiv", 20 * 10000),
-    ("dataset/pretrain/jp", 17 * 10000),
-    ("dataset/pretrain/jp_r18", 17 * 10000),
-    ("dataset/pretrain/jp_r18_rpgmaker", 6 * 10000),
+    ("dataset/pretrain/jp", 15 * 10000),
+    ("dataset/pretrain/jp_r18", 15 * 10000),
+    ("dataset/pretrain/jp_r18_rpgmaker", 10 * 10000),
+    ("dataset/pretrain/kr", 40 * 10000),
 ]
 
 # 加载分词器
@@ -175,11 +176,9 @@ def load_dataset(tokenizer):
             count = count + 1
             with open(f"{dir_path}/{MODEL_NAME}_{dir_name}.txt", "r", encoding = "utf-8") as file:
                 datas_by_type = [line.strip() for line in tqdm(file, desc = path, total = num)]
-                random.shuffle(datas_by_type)
         else:
-            total = len([entry for entry in os.listdir(path) if os.path.isfile(os.path.join(path, entry))])
-
             lines = []
+            total = len([f for f in os.scandir(path) if f.name.endswith(".txt")])
             for file in tqdm(os.scandir(path), desc = path, total = total):
                 if file.name.endswith(".txt"):
                     with open(file.path, "r", encoding = "utf-8") as file:
@@ -201,6 +200,7 @@ def load_dataset(tokenizer):
         datas.extend(datas_by_type)
 
     # 生成数据集
+    random.shuffle(datas)
     os.makedirs("cache", exist_ok = True)
     dataset_train = Dataset.from_dict({"line": datas})
     dataset_train_tokenized = dataset_train.map(

02.py

@@ -6,185 +6,34 @@
 
 from datetime import datetime
 
-import cohere
 from rich import print
 from openai import AsyncOpenAI
 from aiolimiter import AsyncLimiter
 
-PROMPT_JP = (
-"""
-请生成用于实体识别模型训练的日文合成语料，并检查其质量。
-生成时，请遵循以下内容要求、实体类别和质量标准：
+# 设置任务参数
+BATCH = 8
+TIMEOUT = 180
+MAX_LOOP = 32
+LOOP_SIZE = 128
+
+MODEL = "glm-4-9b-chat"
+API_KEY = "sk-no-key-required"
+BASE_URL = "http://localhost:8080/v1"
+TOP_P = 0.95
+TEMPERATURE = 0.95
+PRESENCE_PENALTY = 0.95
+FREQUENCY_PENALTY = 0.00
+
+PROMPT = {}
+for f in os.scandir(f"prompt"):
+    if f.name.startswith("llm_corpus"):
+        with open(f"prompt/{f.name}", "r", encoding = "utf-8") as file:
+            PROMPT[f.name.replace(".txt", "")] = file.read().strip()
+
+LANGUAGE = "kr"
+OUTPUT_PATH = f"dataset/ner/{LANGUAGE}"
+TARGET_PROMPT = PROMPT.get(f"llm_corpus_{LANGUAGE}")
 
-内容要求：
-1、生成语句：
-生成10个语句，每个语句包含2-4个不同类别的实体。
-每个类别的实体在每个语句中最多出现一次，以确保多样性。
-
-2、实体使用：
-每个语句中的实体词语之间不要相互包含。
-实体词语应为日文片假名或平假名形式，避免使用英文或汉字。
-
-3、符号使用：
-除语法上必要的情况外，避免使用《》、「」、『』等符号包裹实体词语。
-
-4、多样性与独特性：
-语句应展现多样性，避免重复或相似的语句结构和实体。
-使用随机性和不同的句子模板来增加多样性。
-
-5、语句类型：
-语句类型应包括但不限于旁白、对话、场景描述、第一人称视角、第三人称视角等。
-
-6、题材涵盖：
-语句题材应涵盖异世界、转生、穿越、奇幻、冒险、战争、科幻、历史、战国、中华风、中世纪、超能力、校园恋爱、运动竞技等轻小说常见题材。
-
-实体类别：
-1、人名（PER）：包括个体的人名，常见的名字、昵称、艺名、历史人物名字等，不包括代指人的称谓、头衔、职业和代词等。
-2、组织与团体（ORG）：包括公司、机构、政府组织、非政府组织、学校、家族、门派等组织与团体。
-3、地点与设施（LOC）：包括国家、城市、州、省、街道、自然地理实体（如河流、山脉）等地点或建筑物、地标、机场、桥梁、剧院、体育场等设施。
-4、产品与道具（PRD）：包括物品、道具、商品、品牌、技术产品等。
-5、事件（EVT）：包括历史事件、会议、发布会、庆典、比赛等。
-
-
-质量标准：
-1、生成的句子应具备高语言质量，确保流畅且自然。
-2、各类实体在句子中的分布应合理，避免单一类型实体的过多重复。
-3、确保生成的句子在日文语境中具有逻辑性和可读性，避免语法错误或不自然的表达。
-4、验证实体的使用是否符合其定义，确保它们在上下文中扮演合理的角色。
-
-回复使用JSON格式，回复中仅需要以下数据，不要出现其他文字或者描述：
-[
-    {
-        "sentence": "<日文语句>",
-        "entities": [
-            {"name": "<实体名称>", "ner_type": "<PER/ORG/LOC/INS/PRD/EVT>"},
-            {"name": "<实体名称>", "ner_type": "<PER/ORG/LOC/INS/PRD/EVT>"}
-        ]
-    }
-]
-"""
-)
-
-PROMPT_CN = (
-"""
-请生成用于实体识别模型训练的中文合成语料，并检查其质量。
-生成时，请遵循以下内容要求、实体类别和质量标准：
-
-内容要求：
-1、生成语句：
-生成10个语句，每个语句包含2-4个不同类别的实体。
-每个类别的实体在每个语句中最多出现一次，以确保多样性。
-
-2、实体使用：
-每个语句中的实体词语之间不要相互包含。
-
-3、符号使用：
-除语法上必要的情况外，避免使用《》、「」、『』等符号包裹实体词语。
-
-4、多样性与独特性：
-语句应展现多样性，避免重复或相似的语句结构和实体。
-使用随机性和不同的句子模板来增加多样性。
-
-5、语句类型：
-语句类型应包括但不限于旁白、对话、场景描述、第一人称视角、第三人称视角等。
-
-6、题材涵盖：
-语句题材应涵盖异世界、转生、穿越、奇幻、冒险、战争、科幻、历史、中世纪、超能力、校园恋爱、运动竞技等轻小说常见题材。
-
-实体类别：
-1、人名（PER）：包括个体的人名，常见的名字、昵称、艺名、历史人物名字等，不包括代指人的称谓、头衔、职业和代词等。
-2、组织与团体（ORG）：包括公司、机构、政府组织、非政府组织、学校、家族、门派等组织与团体。
-3、地点与设施（LOC）：包括国家、城市、州、省、街道、自然地理实体（如河流、山脉）等地点或建筑物、地标、机场、桥梁、剧院、体育场等设施。
-4、产品与道具（PRD）：包括物品、道具、商品、品牌、技术产品等。
-5、事件（EVT）：包括历史事件、会议、发布会、庆典、比赛等。
-
-质量标准：
-1、生成的句子应具备高语言质量，确保流畅且自然。
-2、各类实体在句子中的分布应合理，避免单一类型实体的过多重复。
-3、确保生成的句子应具有逻辑性和可读性，避免语法错误或不自然的表达。
-4、验证实体的使用是否符合其定义，确保它们在上下文中扮演合理的角色。
-
-回复使用JSON格式，回复中仅需要以下数据，不要出现其他文字或者描述：
-[
-    {
-        "sentence": "<中文句子>",
-        "entities": [
-            {"name": "<实体名称>", "ner_type": "<PER/ORG/LOC/PRD/EVT>"},
-            {"name": "<实体名称>", "ner_type": "<PER/ORG/LOC/PRD/EVT>"}
-        ]
-    }
-]
-"""
-)
-
-PROMPT_EN = (
-"""
-请生成用于实体识别模型训练的英文合成语料，并检查其质量。
-生成时，请遵循以下内容要求、实体类别和质量标准：
-
-内容要求：
-1、生成语句：
-生成10个语句，每个语句包含2-4个不同类别的实体。
-每个类别的实体在每个语句中最多出现一次，以确保多样性。
-
-2、实体使用：
-每个语句中的实体词语之间不要相互包含。
-
-3、符号使用：
-除语法上必要的情况外，避免使用《》、「」、『』等符号包裹实体词语。
-
-4、多样性与独特性：
-语句应展现多样性，避免重复或相似的语句结构和实体。
-使用随机性和不同的句子模板来增加多样性。
-
-5、语句类型：
-语句类型应包括但不限于旁白、对话、场景描述、第一人称视角、第三人称视角等。
-
-6、题材涵盖：
-语句题材应涵盖异世界、转生、穿越、奇幻、冒险、战争、科幻、历史、中世纪、超能力、校园恋爱、运动竞技等轻小说常见题材。
-
-实体类别：
-1、人名（PER）：包括个体的人名，常见的名字、昵称、艺名、历史人物名字等，不包括代指人的称谓、头衔、职业和代词等。
-2、组织与团体（ORG）：包括公司、机构、政府组织、非政府组织、学校、家族、门派等组织与团体。
-3、地点与设施（LOC）：包括国家、城市、州、省、街道、自然地理实体（如河流、山脉）等地点或建筑物、地标、机场、桥梁、剧院、体育场等设施。
-4、产品与道具（PRD）：包括物品、道具、商品、品牌、技术产品等。
-5、事件（EVT）：包括历史事件、会议、发布会、庆典、比赛等。
-
-质量标准：
-1、生成的句子应具备高语言质量，确保流畅且自然。
-2、各类实体在句子中的分布应合理，避免单一类型实体的过多重复。
-3、确保生成的句子应具有逻辑性和可读性，避免语法错误或不自然的表达。
-4、验证实体的使用是否符合其定义，确保它们在上下文中扮演合理的角色。
-
-回复使用JSON格式，回复中仅需要以下数据，不要出现其他文字或者描述：
-[
-    {
-        "sentence": "<英文句子>",
-        "entities": [
-            {"name": "<实体名称>", "ner_type": "<PER/ORG/LOC/PRD/EVT>"},
-            {"name": "<实体名称>", "ner_type": "<PER/ORG/LOC/PRD/EVT>"}
-        ]
-    }
-]
-"""
-)
-
-with open("02.json", "r", encoding = "utf-8") as f:
-    data = json.load(f)
-    MODEL = data.get("MODEL", "glm-4-9b-chat")
-    API_KEY = data.get("API_KEY", "sk-no-key-required")
-    BASE_URL = data.get("BASE_URL", "http://localhost:8080/v1")
-
-BATCH = 32
-TIMEOUT = 120
-MAX_LOOP = 3
-LOOP_SIZE = 256
-TEMPERATURE = 1.25
-PRESENCE_PENALTY = 1.0
-
-PROMPT = PROMPT_CN
-
-names = set()
 semaphore = asyncio.Semaphore(BATCH)
 async_limiter = AsyncLimiter(max_rate = BATCH, time_period = 1)
 openai_handler = AsyncOpenAI(
@@ -223,19 +72,20 @@ def fix_broken_json_string(jsonstring):
 
     return jsonstring
 
+# 异步请求
 async def request():
     async with semaphore, async_limiter:
         completion = await openai_handler.chat.completions.create(
             model = MODEL,
             temperature = TEMPERATURE,
-            top_p = 0.5,
-            # max_tokens = 4096,
+            top_p = TOP_P,
+            max_tokens = 3 * 1024,
             presence_penalty = PRESENCE_PENALTY,
-            frequency_penalty = 0,
+            frequency_penalty = FREQUENCY_PENALTY,
             messages = [
                 {
                     "role": "user", 
-                    "content": PROMPT.replace("{words}", ",".join(names))
+                    "content": TARGET_PROMPT
                 },
             ],
         )
@@ -252,13 +102,21 @@ async def request():
             print(message.content.strip())
             raise e
 
-        for v1 in result:
-            print(v1)
-            for v2 in v1.get("entities", []):
-                names.add(v2["name"])
+        for v in result:
+            v["sentence"] = v["sentence"].replace("（PER）", "")
+            v["sentence"] = v["sentence"].replace("（ORG）", "")
+            v["sentence"] = v["sentence"].replace("（LOC）", "")
+            v["sentence"] = v["sentence"].replace("（PRD）", "")
+            v["sentence"] = v["sentence"].replace("（EVT）", "")
+            v["sentence"] = v["sentence"].replace("「", "")
+            v["sentence"] = v["sentence"].replace("」", "")
+            v["sentence"] = v["sentence"].replace("『", "")
+            v["sentence"] = v["sentence"].replace("』", "")
+            print(f"{v}\n")
 
         return result
 
+# 异步任务完成回调
 def on_task_done(future, datas, loop, failed, successed):
     try:
         data = future.result()
@@ -270,13 +128,13 @@ def on_task_done(future, datas, loop, failed, successed):
     finally:
         print(f"正在进行第 {loop} 轮任务，成功 {len(successed)} 次 ... 失败 {len(failed)} 次 ...")
 
+# 主函数
 async def main():
     loop = 0
     start_time = datetime.now().strftime("%Y%m%d_%H%M%S")
 
     while loop < MAX_LOOP:
         loop = loop + 1
-        names = set()
         failed = []
         successed = []
 
@@ -292,10 +150,11 @@ async def main():
         await asyncio.gather(*tasks, return_exceptions = True)
 
         # 写入本地
-        file_path = f"dataset\\{start_time}_{MODEL.replace("/", "_").replace("-", "_")}_{loop:02d}.json"
+        file_path = f"{OUTPUT_PATH}/{start_time}_{MODEL.replace("/", "_").replace("-", "_")}_{loop:02d}.json"
         with open(file_path, "w", encoding = "utf-8") as file:
             file.write(json.dumps(datas, indent = 4, ensure_ascii = False))
             print(f"第 {loop} 轮已完成，数据已写入 {file_path} ...")
 
+# 入口函数
 if __name__ == "__main__":
     asyncio.run(main())