105 lines
3.6 KiB
Python
105 lines
3.6 KiB
Python
|
|
||
|
|
# Copyright 2024 Alibaba Group Holding Limited. All Rights Reserved.
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
# ==============================================================================
|
||
|
|
|
||
|
|
import json
|
||
|
|
import argparse
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
import torch
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
|
||
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
|
|
|
||
|
|
|
||
|
|
def read_json_field(filename):
|
||
|
|
try:
|
||
|
|
with open(filename, 'r') as file:
|
||
|
|
data = json.load(file)
|
||
|
|
output = []
|
||
|
|
for item in data:
|
||
|
|
instruction = item["instruction"]
|
||
|
|
output = item["output"]
|
||
|
|
output.append({"prompt": instruction, "chosen": output})
|
||
|
|
return output
|
||
|
|
except FileNotFoundError:
|
||
|
|
logging.error("The file was not found.")
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
logging.error("There was an error decoding the JSON file.")
|
||
|
|
except Exception as e:
|
||
|
|
logging.error(f"An error occurred: {e}")
|
||
|
|
|
||
|
|
|
||
|
|
def write_data_to_json_file(data, file_path):
|
||
|
|
try:
|
||
|
|
with open(file_path, 'w') as file:
|
||
|
|
json.dump(data, file, ensure_ascii=False, indent=4)
|
||
|
|
logging.info(f"Data successfully written to {file_path}")
|
||
|
|
except Exception as e:
|
||
|
|
logging.error(f"An error occurred: {e}")
|
||
|
|
|
||
|
|
|
||
|
|
def generate_student_response(data_list, config):
|
||
|
|
# load student model
|
||
|
|
student_tokenizer = AutoTokenizer.from_pretrained(
|
||
|
|
config["models"]["student"],
|
||
|
|
trust_remote_code=True
|
||
|
|
)
|
||
|
|
student_model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
config["models"]["student"],
|
||
|
|
device_map="auto",
|
||
|
|
trust_remote_code=True
|
||
|
|
)
|
||
|
|
outcomes = []
|
||
|
|
for sample in tqdm(data_list, desc="Call remote model and generating responses"):
|
||
|
|
prompt = sample["prompt"]
|
||
|
|
chosen = sample["chosen"]
|
||
|
|
# for student model
|
||
|
|
messages = [
|
||
|
|
{"role": "user", "content": prompt}
|
||
|
|
]
|
||
|
|
text = student_tokenizer.apply_chat_template(
|
||
|
|
messages,
|
||
|
|
tokenize=False,
|
||
|
|
add_generation_prompt=True
|
||
|
|
)
|
||
|
|
model_inputs = student_tokenizer([text], return_tensors="pt").to(student_model.device)
|
||
|
|
generated_ids = student_model.generate(
|
||
|
|
**model_inputs,
|
||
|
|
max_new_tokens=config["inference"]["max_new_tokens"]
|
||
|
|
)
|
||
|
|
generated_ids = [
|
||
|
|
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
|
||
|
|
]
|
||
|
|
|
||
|
|
rejected = student_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
||
|
|
gen_data = {'prompt': prompt, 'chosen': chosen, 'rejected': rejected}
|
||
|
|
outcomes.append(gen_data)
|
||
|
|
write_data_to_json_file(outcomes, config["dataset"]["labeled_path"])
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument('--config', type=str, required=True, help='path to the json config file')
|
||
|
|
args = parser.parse_args()
|
||
|
|
config = json.load(open(args.config))
|
||
|
|
data_list = read_json_field(config["dataset"]["instruction_path"])
|
||
|
|
generate_student_response(data_list, config)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|