181 lines
6.4 KiB
Python
181 lines
6.4 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 re
|
||
from tqdm import tqdm
|
||
from openai import OpenAI
|
||
from collections import Counter
|
||
|
||
|
||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
|
||
|
||
predefined_distribution = {
|
||
'Math': 0.167,
|
||
'Code Generation': 0.083,
|
||
'Writing': 0.017,
|
||
'Computer Science': 0.017,
|
||
'Reasoning': 0.167,
|
||
'Complex Format': 0.017,
|
||
'Code Debug': 0.083,
|
||
'Common-Sense': 0.017,
|
||
'Counterfactual': 0.017,
|
||
'Multilingual': 0.017,
|
||
'Roleplay': 0.017,
|
||
'Biology': 0.017,
|
||
'Technology': 0.017,
|
||
'Ethics': 0.017,
|
||
'Sport': 0.017,
|
||
'Law': 0.017,
|
||
'Medicine': 0.017,
|
||
'Literature': 0.017,
|
||
'Entertainment': 0.017,
|
||
'Art': 0.017,
|
||
'Music': 0.017,
|
||
'Toxicity': 0.017,
|
||
'Economy': 0.017,
|
||
'Physics': 0.017,
|
||
'History': 0.017,
|
||
'Chemistry': 0.017,
|
||
'Philosophy': 0.017,
|
||
'Health': 0.017,
|
||
'Ecology': 0.017,
|
||
'Grammar': 0.017,
|
||
'Paraphrase': 0.017,
|
||
'Others': 0.041
|
||
}
|
||
|
||
predefined_prompt = """
|
||
You are a data annotation expert. Please classify the task type or domain of #Given Instruction.
|
||
The task type or domain should be in the list: [’Math’, ’Code Generation’, ’Writing’, ’Computer Science’,
|
||
’Reasoning’, ’Complex Format’, ’Code Debug’, ’Common-Sense’, ’Counterfactual’, ’Multilingual’, ’Roleplay’,
|
||
’Biology’, ’Technology’, ’Ethics’, ’Sport’, ’Law’, ’Medicine’, ’Literature’, ’Entertainment’, ’Art’, ’Music’,
|
||
’Toxicity’, ’Economy’, ’Physics’, ’History’, ’Chemistry’, ’Philosophy’,’Health’,’Ecology’,’Grammar’,’Paraphrase’,
|
||
’Others’]. You should place your answer enclosed within <answer></answer> tags, such as <answer>Math</answer>.
|
||
Do not return anything else.
|
||
#Given Instruction#:
|
||
"""
|
||
|
||
def extract_answer(content):
|
||
pattern = r'<answer>(.*?)</answer>'
|
||
match = re.search(pattern, content, re.DOTALL)
|
||
if match:
|
||
return match.group(1)
|
||
else:
|
||
return None
|
||
|
||
|
||
def read_json_fields(filename):
|
||
try:
|
||
with open(filename, 'r') as file:
|
||
data = json.load(file)
|
||
return data
|
||
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 classify_instruction(instruction, client, model, config):
|
||
message = [
|
||
{"role": "user", "content": predefined_prompt + "\n" + instruction}
|
||
]
|
||
completion = client.chat.completions.create(
|
||
messages = message,
|
||
model = model,
|
||
max_completion_tokens = config["inference"]["max_new_tokens"]
|
||
)
|
||
result = completion.choices[0].message.content.strip()
|
||
result = extract_answer(result)
|
||
if result is None or result not in predefined_distribution.keys():
|
||
result = 'Others'
|
||
return result
|
||
|
||
|
||
def generate_teacher_response_api(data_list, config):
|
||
client = OpenAI(
|
||
api_key = config["inference"]["api_key"],
|
||
base_url = config["inference"]["base_url"]
|
||
)
|
||
models = client.models.list()
|
||
model = models.data[0].id
|
||
logging.info(model)
|
||
|
||
classified_data = []
|
||
for sample in tqdm(data_list, desc="Call remote model and generating responses"):
|
||
instruction = sample["instruction"]
|
||
category = classify_instruction(item['instruction'], client, model)
|
||
new_sample = sample.copy()
|
||
new_sample['category'] = category
|
||
classified_data.append(new_sample)
|
||
|
||
# Count occurrences per category
|
||
category_counts = Counter(item['category'] for item in classified_data)
|
||
total_samples = len(classified_data)
|
||
|
||
# Resample according to predefined distribution
|
||
resampled_data = []
|
||
for category, target_ratio in predefined_distribution.items():
|
||
target_count = int(total_samples * target_ratio)
|
||
category_samples = [item for item in classified_data if item['category'] == category]
|
||
if len(category_samples) == 0:
|
||
logging.warning("No instructions are provided for the category: " + category)
|
||
continue
|
||
if len(category_samples) > target_count:
|
||
# Randomly sample the required number of instructions
|
||
resampled_category_samples = random.sample(category_samples, target_count)
|
||
else:
|
||
# If not enough samples, repeat the existing ones
|
||
resampled_category_samples = category_samples * (target_count // len(category_samples)) + random.sample(
|
||
category_samples, target_count % len(category_samples))
|
||
resampled_data.extend(resampled_category_samples)
|
||
|
||
write_data_to_json_file(resampled_data, config["dataset"]["output_path"])
|
||
|
||
|
||
def sample_with_teacher_model(config):
|
||
logging.info('Generating distillation data from the teacher model!')
|
||
data_list = read_json_fields(config["dataset"]["input_path"])
|
||
job_type = config["job_type"]
|
||
assert job_type == "instruct_sample_api"
|
||
generate_teacher_response_api(data_list, config)
|
||
|
||
|
||
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))
|
||
sample_with_teacher_model(config)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |