1 line
23 KiB
JSON
1 line
23 KiB
JSON
{"repo": "binary-husky/hmp2g", "files": {"cradle.py": {"607dfd213e7490e67813860d30f313b3544d3dc8": {"t": 1785768949.4539123, "n": 3125, "h": "ee91c66feecf759d", "c": "\nclass GptJsonIO():\n def __init__(self, schema):\n self.schema = schema\n self.prompt = \"You should only respond in JSON format as described below.\\nResponse Format:\\n{json_schema}\\nEnsure the response can be parsed by Python json.loads.\"\n\n def generate_input(self):\n from gpt_json import GPTJSON, GPTMessage, GPTMessageRole\n schema, prompt = self.schema, self.prompt\n self.gptjson = GPTJSON[schema](\"^_^\")\n\n messages=[\n GPTMessage(\n role=GPTMessageRole.SYSTEM,\n content=prompt,\n ),\n ]\n\n format_variables = None\n truncation_options = None\n max_response_tokens = None\n messages = self.gptjson.fill_messages(\n messages, format_variables, truncation_options, max_response_tokens\n )\n return messages[0].content\n\n def generate_output(self, response):\n _response = {\"choices\":[{\"message\":{\"content\":response}}]}\n extracted_json, fixed_payload = self.gptjson.extract_json(_response, self.gptjson.extract_type)\n return self.gptjson.schema_model(**extracted_json)\n\n def generate_repair_prompt(self, broken_json):\n import json\n error = \"\"\n try:\n json.loads(broken_json)\n except Exception as e:\n error = e.args[0]\n prompt = f\"The following json string is broken, error message is \\n```\\n{error}\\n```\\nNow repair it: \\n\\n\" + broken_json\n return prompt\n \n def generate_output_auto_repair(self, response, gpt_gen_fn):\n \"\"\"\n response: string containing canidate json\n gpt_gen_fn: gpt_gen_fn(inputs, sys_prompt)\n \"\"\"\n try:\n result = self.generate_output(response)\n except:\n try:\n # json 格式异常，尝试修复一次\n print('Repairing json：', response)\n repair_prompt = gjio.generate_repair_prompt(broken_json = response)\n result = gjio.generate_output(gpt_gen_fn(repair_prompt, formatting_sys_prompt))\n print('Repaire json success!', result)\n except Exception as e:\n # 没辙了，放弃治疗\n print('Repaire json fail!')\n raise RuntimeError('Cannot repair json.', str(e))\n return result\n\nimport void_terminal as vt\nfrom void_terminal.request_llm.bridge_all import predict_no_ui_long_connection\nvt.set_conf(key=\"API_KEY\", value=\"fk195831-IdP9Pb3W6DCMUIbQwVX6MsSiyxwqybyS\")\nvt.set_conf(key=\"LLM_MODEL\", value=\"api2d-gpt-3.5-turbo\")\n\nfrom pydantic import BaseModel, Field\nclass Schema(BaseModel):\n revised_answer: str = Field(description=\"the revised answer.\")\ngjio = GptJsonIO(Schema)\nformatting_sys_prompt = gjio.generate_input()\n# res = gjio.generate_output('{\\n\"revised_answer\": \"两者都可以，宣传可以推动社会进步，而推动社会进步也可以宣传。\",\\n}')\nres = '{\\n\"revised_answer\": \"两者都可以，\"宣传\"可以推动社会进步，而推动社会进步也可以宣传。\",\\n}'\nres = gjio.generate_output_auto_repair(\n res, \n lambda x, p: predict_no_ui_long_connection(inputs=x, llm_kwargs=vt.get_chat_default_kwargs()['llm_kwargs'], history=[], sys_prompt=p)\n)\nprint(res.revised_answer)\n\n"}}, "ALGORITHM/multi_llm/foundation.py": {"607dfd213e7490e67813860d30f313b3544d3dc8": {"t": 1785768958.1104634, "n": 18799, "h": "d8142456056f5664", "c": "import os, time, torch, json, random\nimport numpy as np\nfrom UTIL.colorful import *\nfrom config import GlobalConfig\nfrom UTIL.tensor_ops import repeat_at\nfrom ALGORITHM.common.rl_alg_base import RLAlgorithmBase\nfrom ALGORITHM.common.logit2act import Logit2Act\nfrom .temp import sample_good_QA_from_turns, get_log_prob, get_log_probs_with_input_ids\nfrom UTIL.tensor_ops import repeat_at, _2tensor, scatter_righthand, my_view\nfrom .temp import RewardBySimilarity, gae_vectorize\nfrom .bridge_llm import ChatGLMCritic\nfrom pathlib import Path\nfrom torch.optim import Adam\nfrom .temp import tokenize_qa\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport void_terminal as vt\nfrom .llm_util_base import VtTopicLauncher, TopicLauncher, RewardEvaluator, Reviser, SimpleDatasetManager\n\n# vt.set_conf(key=\"API_KEY\", value=\"6424e9d19e674092815cea1cb35e67a5\")\n# vt.set_conf(key=\"AZURE_ENDPOINT\", value=\"https://rhtjjjjjj.openai.azure.com/\")\n# vt.set_conf(key=\"AZURE_ENGINE\", value=\"qqwe\")\n# vt.set_conf(key=\"LLM_MODEL\", value=\"azure-gpt-3.5\")\n\nvt.set_conf(key=\"API_KEY\", value=\"fk195831-IdP9Pb3W6DCMUIbQwVX6MsSiyxwqybyS\")\nvt.set_conf(key=\"LLM_MODEL\", value=\"api2d-gpt-3.5-turbo\")\n\nclass Generator:\n def __init__(self, gen):\n self.gen = gen\n\n def __iter__(self):\n self.value = yield from self.gen\n\nclass AlgorithmConfig:\n '''\n AlgorithmConfig: This config class will be 'injected' with new settings from json.\n (E.g., override configs with ```python main.py --cfg example.jsonc```)\n (please see UTIL.config_args to find out how this advanced trick works out.)\n '''\n # configuration, open to jsonc modification\n gamma = 0.99\n tau = 0.95\n train_traj_needed = 512\n dataset_path = './ALGORITHM/llm/profile_instance.json.insert.json'\n model_path = \"/home/hmp/Miraclemarvel55_RLHF/chatglm_6b\"\n max_source_length = 100\n max_gen_tokens = 100\n mini_batch_size = 16\n sft_batch_size = 16\n debug = False\n\n # device_override = \"cuda:1\"\n # device_launcher_llm = 'cuda:1'\n # device_main_llm = 'cuda:1'\n # device_reward_llm = 'cuda:1'\n\n device_override = \"cuda:0,1\"\n device_launcher_llm = 'cuda:0'\n device_main_llm = 'cuda:1'\n device_reward_llm = 'cuda:0'\n\n enable_rl = False\n enable_stf = True\n\n save_step = 50\n topics = [\n \"房价过高\", \"内卷\", \"AI威胁\", \"量子计算\",\n \"房地产调控\", \"汇率\", \"字幕组\", \"网红\",\n \"GDP\", \"家庭\", \"视觉\", \"守望先锋\",\n \"金砖国家\",\"俄罗斯\",\"月球\",\"美国中央情报局\",\n \"被骗到缅北\", \"平壤\", \"特斯拉\", \"未来生活\"\n ]\n topic_launcher_prompt = f\"针对“TOPIC_REPLACE”话题，提出一个引人深思的问题。限20字。\"\n reward_model_prompt = f\"问题：Q_REPLACE\\n\\n回答：A_REPLACE\\n\\n判断：以上文本的回答部分是否通俗，让中学学历的人也能看懂？仅回答“是”或“否”。\"\n # reviser_prompt = 'Answer with Chinese, and you must use a metaphor and give a short answer. ' # 这个prompt建议使用英文\n reviser_prompt_input = \"\\n使用抬杠的语气否定该问题的正确性。将你的回复放入json的answer字段，40字以内。\"\n\n\n dataset_size = 1000\n\ndef override_cuda_settings(AlgorithmConfig):\n # change Local cuda settings according to AlgorithmConfig\n if AlgorithmConfig.device_override != \"no-override\":\n assert GlobalConfig.device == 'cpu', \"please set GlobalConfig.device=cpu if you want to use different GPUs for different teams\"\n # reflesh the cuda setting inherited from main.py\n GlobalConfig.device = AlgorithmConfig.device_override\n from main import pytorch_gpu_init; pytorch_gpu_init(GlobalConfig)\n # reflesh the cached cuda setting in tensor_ops\n from UTIL.tensor_ops import cuda_cfg; cuda_cfg.read_cfg()\n\nclass LLM_Foundation(RLAlgorithmBase, VtTopicLauncher, RewardEvaluator, Reviser, SimpleDatasetManager):\n\n def __init__(self, n_agent, n_thread, space, mcv=None, team=None):\n override_cuda_settings(AlgorithmConfig)\n self.n_agent = n_agent\n self.topics = AlgorithmConfig.topics\n from .bridge_llm import load_static_model, load_trainable_model, load_trainable_headless_model\n # 主模型\n self.main_llm_model, self.tokenizer = load_trainable_model(device=AlgorithmConfig.device_main_llm)\n self.save_model(out_dir=f'RESULT/{GlobalConfig.note}/llm_model')\n # 高温模型\n self.launcher_llm_model, self.tokenizer = load_static_model(device=AlgorithmConfig.device_launcher_llm)\n # # 奖励模型 （效果太差，使用更强的gpt-3.5模型）\n # self.reward_llm_model, self.tokenizer = self.launcher_llm_model, self.tokenizer # load_static_model(device=AlgorithmConfig.device_reward_llm)\n\n self.critic = ChatGLMCritic(device=AlgorithmConfig.device_main_llm)\n optimize_params = self.main_llm_model.training_para_list + list(self.critic.parameters())\n self.optimizer = Adam(optimize_params, lr=1e-4, eps=1e-3)\n self.logic_processor = Logit2Act()\n\n self.mcv = mcv\n\n \n def main_llm_answer_question(self, topic_launcher_questions):\n def pad(arr):\n pad_id = self.tokenizer.convert_tokens_to_ids(\"<pad>\")\n len_list = [len(token_qa_prompt) for i, token_qa_prompt in enumerate(arr)]\n max_len = max(len_list)\n for i, token_qa_prompt in enumerate(arr):\n arr[i] = [pad_id] * (max_len - len_list[i]) + token_qa_prompt\n return np.array(arr)\n token_q_prompt_array = []\n for q in topic_launcher_questions:\n token_q_prompt, gen_len, prompt = tokenize_qa(self.tokenizer, query=q, history=[])\n token_q_prompt = token_q_prompt['input_ids'].tolist()[0]\n token_q_prompt_array.append(token_q_prompt)\n token_q_prompt_array = pad(token_q_prompt_array)\n\n input_ids = _2tensor(token_q_prompt_array).to(AlgorithmConfig.device_main_llm)\n tokenlen_q_pad = question_token_len = input_ids.shape[-1]\n num_beams = num_return_sequences = 1\n assert num_beams >= num_return_sequences, \"candidates num should greater than returns num\"\n # max_new_tokens = 8\n gen_method = \"greedy_search\" # if num_beams == 1 else \"beam_search\" \n main_model_result = self.main_llm_model.generate(input_ids=input_ids, do_sample=False, num_beams=num_beams, max_new_tokens=AlgorithmConfig.max_gen_tokens,\n num_return_sequences=num_return_sequences, use_cache=True, num_beam_groups=1, output_scores=True,\n output_hidden_states=False, return_dict_in_generate=True)\n main_sequences = main_model_result.sequences\n main_log_probs = get_log_prob(generated_outputs=main_model_result, input_ids=input_ids, gen_method=gen_method)\n main_gen_texts = self.tokenizer.batch_decode(main_sequences)\n main_gen_texts_only_answer = self.tokenizer.batch_decode(main_sequences[:,question_token_len:])\n for c in main_gen_texts_only_answer: print(c)\n return main_sequences, main_log_probs, main_gen_texts_only_answer, tokenlen_q_pad\n\n\n\n def main_llm_sft(self, n_step):\n from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler\n\n dataset = self.read_ds(tag='sft')\n\n\n for _ in range(n_step):\n big_batch_size = len(dataset)\n self.sampler = BatchSampler(SubsetRandomSampler(range(big_batch_size)), AlgorithmConfig.sft_batch_size, drop_last=False)\n self.sampler = iter(self.sampler)\n indices = next(self.sampler)\n samples = dataset[indices]\n\n def pad(arr):\n pad_id = self.tokenizer.convert_tokens_to_ids(\"<pad>\")\n len_list = [len(token_qa_prompt) for i, token_qa_prompt in enumerate(arr)]\n max_len = max(len_list)\n for i, token_qa_prompt in enumerate(arr):\n arr[i] = [pad_id] * (max_len - len_list[i]) + token_qa_prompt\n return np.array(arr)\n \n token_q_prompt_array = []\n for q, a in samples:\n a = a[:AlgorithmConfig.max_gen_tokens]\n token_q_prompt, gen_len, prompt = tokenize_qa(self.tokenizer, query=\"\", history=[[q, a]])\n token_q_prompt = token_q_prompt['input_ids'].tolist()[0]\n token_q_prompt_array.append(token_q_prompt)\n token_q_prompt_array = pad(token_q_prompt_array)\n\n input_ids = _2tensor(token_q_prompt_array).to(AlgorithmConfig.device_main_llm)\n tokenlen_q_pad = question_token_len = input_ids.shape[-1]\n num_beams = num_return_sequences = 1\n assert num_beams >= num_return_sequences, \"candidates num should greater than returns num\"\n\n # max_new_tokens = 8\n model_inputs = self.main_llm_model.prepare_inputs_for_generation(input_ids)\n output = self.main_llm_model(**model_inputs)\n logits = output.logits\n shift_logits = logits[..., :-1, :].contiguous()\n shift_labels = input_ids[..., 1:].contiguous()\n \n loss = nn.functional.cross_entropy(my_view(shift_logits, [-1, 0]), my_view(shift_labels, [-1]))\n\n self.optimizer.zero_grad()\n loss.backward()\n self.optimizer.step()\n\n\n self.mcv.rec(loss.item(),'stf_loss')\n self.mcv.rec_show()\n\n\n print(loss.item())\n\n\n\n\n\n def interact_with_env(self, StateRecall):\n '''\n Interfacing with marl, standard method that you must implement\n (redirect to shell_env to help with history rolling)\n '''\n # from .temp import tokenize_qa\n print('\\n\\n\\n')\n\n if self.ds_size(tag='sft') < AlgorithmConfig.dataset_size:\n # 获取话题问题\n with torch.no_grad(): # 此处不需要梯度，后面PPO会二次计算\n topic_launcher_questions = self.topic_launcher_propose_question()\n torch.cuda.empty_cache()\n print('-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=--=-=-=-=-')\n # \n\n \n # 主模型回答问题\n with torch.no_grad(): # 此处不需要梯度，后面PPO会二次计算\n main_sequences, main_log_probs, main_gen_texts_only_answer, tokenlen_q_pad = self.main_llm_answer_question(topic_launcher_questions)\n torch.cuda.empty_cache()\n print('-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=--=-=-=-=-')\n # \n\n if AlgorithmConfig.enable_rl:\n # 给出奖励\n with torch.no_grad(): # 此处不需要梯度，后面PPO会二次计算\n gen_texts_answer_only, rewards = self.reward_eval_answer_question(topic_launcher_questions, main_gen_texts_only_answer)\n rewards, masks = self.place_reward(tokenlen_q_pad, \n main_sequences, rewards, self.tokenizer.convert_tokens_to_ids(\"<pad>\"))\n torch.cuda.empty_cache()\n print('-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=--=-=-=-=-')\n\n if AlgorithmConfig.enable_stf:\n # 给出修改后的参考答案\n with torch.no_grad(): # 此处不需要梯度，后面PPO会二次计算\n revised_answer = self.revise_answer_question(topic_launcher_questions, main_gen_texts_only_answer)\n torch.cuda.empty_cache()\n print('-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=--=-=-=-=-')\n\n if AlgorithmConfig.enable_stf:\n # 训练\n self.main_llm_sft(5)\n torch.cuda.empty_cache()\n print('-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=--=-=-=-=-')\n\n\n if AlgorithmConfig.enable_rl:\n torch.cuda.empty_cache()\n self.ppo(ppo_epochs=5, sequences=main_sequences, log_probs=main_log_probs, rewards=rewards, masks=masks, clip_param=0.2)\n\n\n action = np.zeros(shape=(1, 1))\n if StateRecall['Current-Obs-Step'][0] % AlgorithmConfig.save_step == 1: # AlgorithmConfig.save_step-1:\n self.save_model(out_dir=f'RESULT/{GlobalConfig.note}/llm_model')\n return action, StateRecall\n \n def flat_samples(self, data_set):\n self.data_set_collection = []\n for i in range(len(data_set)):\n question = data_set[i][0]['question']\n good_ans = data_set[i][0]['good_ans']\n bad_ans = data_set[i][0]['bad_ans']\n self.data_set_collection.append({\n 'question': question,\n 'good_ans_demos': good_ans,\n 'bad_ans_demos': bad_ans,\n 'score': 1\n })\n\n self.data_set_collection = np.array(self.data_set_collection, dtype=object)\n return \n \n\n def sample_qa(self):\n if not hasattr(self, 'sampler'):\n from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler\n big_batch_size = len(self.data_set_collection)\n self.sampler = BatchSampler(SubsetRandomSampler(range(big_batch_size)), AlgorithmConfig.mini_batch_size, drop_last=True)\n self.sampler = iter(self.sampler)\n\n try: \n indices = next(self.sampler)\n except StopIteration:\n from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler\n big_batch_size = len(self.data_set_collection)\n self.sampler = BatchSampler(SubsetRandomSampler(range(big_batch_size)), AlgorithmConfig.mini_batch_size, drop_last=True)\n self.sampler = iter(self.sampler)\n indices = next(self.sampler)\n\n samples = self.data_set_collection[indices]\n pad_id = self.tokenizer.convert_tokens_to_ids(\"<pad>\")\n\n query = [s['question'] for s in samples]\n good_ans_demos = [s['good_ans_demos'] for s in samples]\n bad_ans_demos = [s['bad_ans_demos'] for s in samples]\n\n res = {\n \"token_qa_prompt_array\": [],\n \"token_q_prompt_array\": [],\n \"tokenlen_qa_prompt_array\": [],\n \"tokenlen_q_prompt_array\": [],\n \"good_ans_demos\": [],\n \"bad_ans_demos\": [],\n \"good_ans_demo\": [],\n }\n for i, q in enumerate(query):\n good_ans_demo = random.choice(good_ans_demos[i])\n token_qa_prompt, gen_len, prompt = tokenize_qa(self.tokenizer, query=\"\", history=[[q, good_ans_demo]])\n token_qa_prompt = token_qa_prompt['input_ids'].tolist()[0]\n res[\"token_qa_prompt_array\"].append(token_qa_prompt)\n res[\"tokenlen_qa_prompt_array\"].append(len(token_qa_prompt))\n token_q_prompt, gen_len, prompt = tokenize_qa(self.tokenizer, query=q, history=[])\n token_q_prompt = token_q_prompt['input_ids'].tolist()[0]\n res[\"token_q_prompt_array\"].append(token_q_prompt)\n res[\"tokenlen_q_prompt_array\"].append(len(token_q_prompt))\n res[\"good_ans_demo\"].append(good_ans_demo)\n res[\"good_ans_demos\"].append(good_ans_demos[i])\n res[\"bad_ans_demos\"].append(bad_ans_demos[i])\n\n # pad\n def pad(arr):\n len_list = [len(token_qa_prompt) for i, token_qa_prompt in enumerate(arr)]\n max_len = max(len_list)\n for i, token_qa_prompt in enumerate(arr):\n arr[i] = [pad_id] * (max_len - len_list[i]) + token_qa_prompt\n return np.array(arr)\n \n res[\"token_qa_prompt_array\"] = pad(res[\"token_qa_prompt_array\"])\n res[\"token_q_prompt_array\"] = pad(res[\"token_q_prompt_array\"])\n res[\"tokenlen_q_pad\"] = res[\"token_q_prompt_array\"].shape[-1]\n res[\"tokenlen_qa_pad\"] = res[\"token_qa_prompt_array\"].shape[-1]\n res[\"tokenlen_a_pad\"] = res[\"tokenlen_qa_pad\"] - res[\"tokenlen_q_pad\"]\n return res\n\n def save_model(self, out_dir):\n import shutil\n os.makedirs(out_dir, exist_ok=True)\n model_path = os.path.join(out_dir,'saved')\n shutil.rmtree(model_path, ignore_errors=True)\n # shutil.copytree(AlgorithmConfig.model_path, model_path)\n self.main_llm_model.save_pretrained(model_path)\n self.tokenizer.save_pretrained(model_path)\n\n def place_reward(self, len_of_query, sequences, reward, pad_id):\n reward = _2tensor(reward).to(AlgorithmConfig.device_main_llm)\n rewards = torch.zeros_like( sequences, dtype=reward.dtype, device=reward.device) + np.nan\n rewards[sequences != pad_id] = 0\n rewards[..., :len_of_query] = np.nan\n masks = ~torch.isnan(rewards)\n # push reward forward\n sentence_end_point = (~torch.isnan(rewards)).sum(-1) + len_of_query - 1 # location where we should insearch the reward\n rewards = scatter_righthand(rewards, reward, sentence_end_point.unsqueeze(-1), check=True) # insert reward\n rewards = torch.nan_to_num(rewards, 0)\n return rewards, masks\n\n\n def ppo(self, ppo_epochs, sequences, log_probs, rewards, masks, clip_param):\n for ppo_epoch in range(ppo_epochs):\n # compute new log probs\n new_log_probs = get_log_probs_with_input_ids(self.main_llm_model, sequences, log_probs.shape[1])\n with torch.no_grad():\n new_log_probs_ref = None # get_log_probs_with_input_ids(self.main_llm_model_ref, sequences.to(self.main_llm_model_ref.device), log_probs.shape[1])\n if new_log_probs_ref is not None:\n kl_div_loss = F.kl_div(new_log_probs_ref.to(new_log_probs.device), new_log_probs, log_target=True, reduction=\"none\").mean()\n \n # entropy = 0 # 暂时不需要熵的约束\n # compute value\n # 到奖励模型和值函数模型的输入可以是一样的都是生成的序列。\n # 生成序列同时包括state和next action\n # prepare input for critic model\n input_ids_critic = sequences.to(AlgorithmConfig.device_main_llm)\n values = self.critic(input_ids=input_ids_critic)\n # compute gae\n gae = gae_vectorize(values=values, rewards=rewards, masks=masks)\n advantages = gae[:, -log_probs.shape[-1]:].to(new_log_probs.device) # 根据参考输出的长度进行截断\n # 计算value的估计量的偏差作为actor loss\n # 以及ppo的actor_loss\n value_estimator_delta = advantages\n ratio = (new_log_probs - log_probs).exp()\n # print(\"reward\", reward, \"ratio:\", ratio, sep=\"\\n\")\n if torch.isinf(ratio).any():\n break\n surr1 = ratio * advantages\n surr2 = torch.clamp(ratio, 1.0 - clip_param, 1.0 + clip_param) * advantages\n actor_loss = - torch.min(surr1, surr2).mean()\n critic_loss = value_estimator_delta.square().mean()\n\n\n loss = critic_loss + actor_loss\n if new_log_probs_ref is not None:\n loss += kl_div_loss * 0.001\n\n self.mcv.rec(critic_loss.item(),'critic_loss')\n self.mcv.rec(actor_loss.item(),'actor_loss')\n if new_log_probs_ref is not None:\n self.mcv.rec(kl_div_loss.item(),'kl_div_loss')\n self.mcv.rec_show()\n # optimize\n self.optimizer.zero_grad()\n loss.backward()\n self.optimizer.step()\n print(\"loss\", loss.item())\n"}}}} |