274 lines
11 KiB
Python
274 lines
11 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 jsonlines
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
from jinja2 import Environment, FileSystemLoader
|
||
|
|
from vllm import LLM, SamplingParams
|
||
|
|
from tqdm import tqdm
|
||
|
|
from openai import OpenAI
|
||
|
|
|
||
|
|
from utils import write_data_to_json_file
|
||
|
|
|
||
|
|
|
||
|
|
# I have checked this function.
|
||
|
|
def cot_generate_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
|
||
|
|
prompt = config["inference"]["prompt"]
|
||
|
|
stream = config["inference"]["stream"]
|
||
|
|
logging.info(model)
|
||
|
|
outcomes = []
|
||
|
|
for sample in tqdm(data_list, desc="Calling remote model and generating responses"):
|
||
|
|
sample = prompt + "\n" + sample
|
||
|
|
message = [
|
||
|
|
{"role": "user", "content": sample}
|
||
|
|
]
|
||
|
|
completion = client.chat.completions.create(
|
||
|
|
messages = message,
|
||
|
|
model = model,
|
||
|
|
max_completion_tokens = config["inference"]["max_new_tokens"],
|
||
|
|
stream = stream
|
||
|
|
)
|
||
|
|
if stream:
|
||
|
|
result = ""
|
||
|
|
for chunk in completion:
|
||
|
|
result += chunk.choices[0].delta.content
|
||
|
|
else:
|
||
|
|
result = completion.choices[0].message.content
|
||
|
|
if result is not None:
|
||
|
|
outcomes.append({"instruction": sample, "output": result})
|
||
|
|
write_data_to_json_file(outcomes, config["dataset"]["output_path"])
|
||
|
|
|
||
|
|
|
||
|
|
def cot_generate_batch(tokenizer, llm, data_list, config, batch_size=32):
|
||
|
|
full_path = config["dataset"]["template"]
|
||
|
|
template_dir = os.path.dirname(full_path)
|
||
|
|
template_file = os.path.basename(full_path)
|
||
|
|
env = Environment(loader=FileSystemLoader(template_dir))
|
||
|
|
template = env.get_template(template_file)
|
||
|
|
prompt = config["inference"]["prompt"]
|
||
|
|
|
||
|
|
batches = [data_list[i:i + batch_size] for i in range(0, len(data_list), batch_size)]
|
||
|
|
for batch in tqdm(batches, desc="Generating responses"):
|
||
|
|
new_batch = []
|
||
|
|
for sample in batch:
|
||
|
|
sample = prompt + "\n" + sample
|
||
|
|
logging.info(sample)
|
||
|
|
message={"role": "user", "content": sample}
|
||
|
|
full_text = template.render(
|
||
|
|
message=message,
|
||
|
|
add_generation_prompt=True,
|
||
|
|
add_output=False
|
||
|
|
)
|
||
|
|
new_batch.append(full_text)
|
||
|
|
outputs = llm.generate(
|
||
|
|
new_batch,
|
||
|
|
SamplingParams(
|
||
|
|
n=1,
|
||
|
|
top_k=1,
|
||
|
|
temperature=config["inference"]["temperature"],
|
||
|
|
seed=config["inference"]["seed"],
|
||
|
|
skip_special_tokens=False,
|
||
|
|
ignore_eos=False,
|
||
|
|
max_tokens=config["inference"]["max_new_tokens"],
|
||
|
|
)
|
||
|
|
)
|
||
|
|
responses = [output.outputs[0].text for output in outputs]
|
||
|
|
outcomes = []
|
||
|
|
for i in range(len(batch)):
|
||
|
|
if responses[i] is not None:
|
||
|
|
outcomes.append((sample,responses[i]))
|
||
|
|
|
||
|
|
with jsonlines.open(config["dataset"]["output_path"], mode='a') as writer:
|
||
|
|
for ins,result in outcomes:
|
||
|
|
gen_data = {"instruction": ins, "output": result}
|
||
|
|
writer.write(gen_data)
|
||
|
|
|
||
|
|
|
||
|
|
def cot_long2short_api(data_list_ins, data_list_out, config):
|
||
|
|
client = OpenAI(
|
||
|
|
api_key = config["inference"]["api_key"],
|
||
|
|
base_url = config["inference"]["base_url"],
|
||
|
|
)
|
||
|
|
models = client.models.list()
|
||
|
|
model = models.data[0].id
|
||
|
|
prompt = config["inference"]["prompt"]
|
||
|
|
stream = config["inference"]["stream"]
|
||
|
|
logging.info(model)
|
||
|
|
outcomes = []
|
||
|
|
data_list=[(ins,out) for ins,out in zip(data_list_ins,data_list_out)]
|
||
|
|
for ins,out in tqdm(data_list, desc="Calling remote model and generating responses"):
|
||
|
|
sample = f"{prompt} Simplify the reasoning process for the problem below.\n\nProblem:\n{ins}\n\nAnswer:\n{out}\n\nSimplified Reasoning Process:"
|
||
|
|
logging.info(sample)
|
||
|
|
message = [
|
||
|
|
{"role": "user", "content": sample}
|
||
|
|
]
|
||
|
|
completion = client.chat.completions.create(
|
||
|
|
messages = message,
|
||
|
|
model = model,
|
||
|
|
max_completion_tokens = config["inference"]["max_new_tokens"],
|
||
|
|
stream = stream,
|
||
|
|
)
|
||
|
|
if stream:
|
||
|
|
result = ""
|
||
|
|
for chunk in completion:
|
||
|
|
result += chunk.choices[0].delta.content
|
||
|
|
else:
|
||
|
|
result = completion.choices[0].message.content
|
||
|
|
|
||
|
|
if result is not None:
|
||
|
|
outcomes.append((sample,result))
|
||
|
|
|
||
|
|
with jsonlines.open(config["dataset"]["output_path"], mode='a') as writer:
|
||
|
|
for ins,result in outcomes:
|
||
|
|
gen_data = {"instruction": ins, "output": result}
|
||
|
|
writer.write(gen_data)
|
||
|
|
|
||
|
|
|
||
|
|
def cot_long2short_batch(tokenizer, llm, data_list_ins, data_list_out, config, batch_size=32):
|
||
|
|
full_path = config["dataset"]["template"]
|
||
|
|
template_dir = os.path.dirname(full_path)
|
||
|
|
template_file = os.path.basename(full_path)
|
||
|
|
env = Environment(loader=FileSystemLoader(template_dir))
|
||
|
|
template = env.get_template(template_file)
|
||
|
|
prompt = config["inference"]["prompt"]
|
||
|
|
data_list=[(ins,out) for ins,out in zip(data_list_ins,data_list_out)]
|
||
|
|
batches = [data_list[i:i + batch_size] for i in range(0, len(data_list), batch_size)]
|
||
|
|
for batch in tqdm(batches, desc="Generating responses"):
|
||
|
|
new_batch = []
|
||
|
|
for ins,out in batch:
|
||
|
|
sample = f"{prompt} Simplify the reasoning process for the problem below.\n\nProblem:\n{ins}\n\nAnswer:\n{out}\n\nSimplified Reasoning Process:"
|
||
|
|
logging.info(sample)
|
||
|
|
message={"role": "user", "content": sample}
|
||
|
|
full_text = template.render(
|
||
|
|
message=message,
|
||
|
|
add_generation_prompt=True,
|
||
|
|
add_output=False
|
||
|
|
)
|
||
|
|
new_batch.append(full_text)
|
||
|
|
outputs = llm.generate(
|
||
|
|
new_batch,
|
||
|
|
SamplingParams(
|
||
|
|
n=1,
|
||
|
|
top_k=1,
|
||
|
|
temperature=config["inference"]["temperature"],
|
||
|
|
seed=config["inference"]["seed"],
|
||
|
|
skip_special_tokens=False,
|
||
|
|
ignore_eos=False,
|
||
|
|
max_tokens=config["inference"]["max_new_tokens"],
|
||
|
|
)
|
||
|
|
)
|
||
|
|
responses = [output.outputs[0].text for output in outputs]
|
||
|
|
outcomes = []
|
||
|
|
for i in range(len(batch)):
|
||
|
|
if responses[i] is not None:
|
||
|
|
outcomes.append((sample,responses[i]))
|
||
|
|
|
||
|
|
with jsonlines.open(config["dataset"]["output_path"], mode='a') as writer:
|
||
|
|
for ins,result in outcomes:
|
||
|
|
gen_data = {"instruction": ins, "output": result}
|
||
|
|
writer.write(gen_data)
|
||
|
|
|
||
|
|
|
||
|
|
def cot_short2long_api(data_list_ins, data_list_out, config):
|
||
|
|
client = OpenAI(
|
||
|
|
api_key = config["inference"]["api_key"],
|
||
|
|
base_url = config["inference"]["base_url"],
|
||
|
|
)
|
||
|
|
models = client.models.list()
|
||
|
|
model = models.data[0].id
|
||
|
|
prompt = config["inference"]["prompt"]
|
||
|
|
stream = config["inference"]["stream"]
|
||
|
|
logging.info(model)
|
||
|
|
outcomes = []
|
||
|
|
data_list=[(ins,out) for ins,out in zip(data_list_ins,data_list_out)]
|
||
|
|
for ins,out in tqdm(data_list, desc="Calling remote model and generating responses"):
|
||
|
|
sample = f"{prompt} Extend the reasoning process for the problem below.\n\nProblem:\n{ins}\n\nAnswer:\n{out}\n\nExtended Reasoning Process:"
|
||
|
|
logging.info(sample)
|
||
|
|
message = [
|
||
|
|
{"role": "user", "content": sample}
|
||
|
|
]
|
||
|
|
completion = client.chat.completions.create(
|
||
|
|
messages = message,
|
||
|
|
model = model,
|
||
|
|
max_completion_tokens = config["inference"]["max_new_tokens"],
|
||
|
|
stream = stream,
|
||
|
|
)
|
||
|
|
if stream:
|
||
|
|
result = ""
|
||
|
|
for chunk in completion:
|
||
|
|
result += chunk.choices[0].delta.content
|
||
|
|
else:
|
||
|
|
result = completion.choices[0].message.content
|
||
|
|
|
||
|
|
if result is not None:
|
||
|
|
outcomes.append((sample,result))
|
||
|
|
|
||
|
|
with jsonlines.open(config["dataset"]["output_path"], mode='a') as writer:
|
||
|
|
for ins,result in outcomes:
|
||
|
|
gen_data = {"instruction": ins, "output": result}
|
||
|
|
writer.write(gen_data)
|
||
|
|
|
||
|
|
|
||
|
|
def cot_short2long_batch(tokenizer, llm, data_list_ins, data_list_out, config, batch_size=32):
|
||
|
|
full_path = config["dataset"]["template"]
|
||
|
|
template_dir = os.path.dirname(full_path)
|
||
|
|
template_file = os.path.basename(full_path)
|
||
|
|
env = Environment(loader=FileSystemLoader(template_dir))
|
||
|
|
template = env.get_template(template_file)
|
||
|
|
prompt = config["inference"]["prompt"]
|
||
|
|
data_list=[(ins,out) for ins,out in zip(data_list_ins,data_list_out)]
|
||
|
|
batches = [data_list[i:i + batch_size] for i in range(0, len(data_list), batch_size)]
|
||
|
|
for batch in tqdm(batches, desc="Generating responses"):
|
||
|
|
new_batch = []
|
||
|
|
for ins,out in batch:
|
||
|
|
sample = f"{prompt} Extend the reasoning process for the problem below.\n\nProblem:\n{ins}\n\nAnswer:\n{out}\n\nExtended Reasoning Process:"
|
||
|
|
logging.info(sample)
|
||
|
|
message={"role": "user", "content": sample}
|
||
|
|
full_text = template.render(
|
||
|
|
message=message,
|
||
|
|
add_generation_prompt=True,
|
||
|
|
add_output=False
|
||
|
|
)
|
||
|
|
new_batch.append(full_text)
|
||
|
|
outputs = llm.generate(
|
||
|
|
new_batch,
|
||
|
|
SamplingParams(
|
||
|
|
n=1,
|
||
|
|
top_k=1,
|
||
|
|
temperature=config["inference"]["temperature"],
|
||
|
|
seed=config["inference"]["seed"],
|
||
|
|
skip_special_tokens=False,
|
||
|
|
ignore_eos=False,
|
||
|
|
max_tokens=config["inference"]["max_new_tokens"],
|
||
|
|
)
|
||
|
|
)
|
||
|
|
responses = [output.outputs[0].text for output in outputs]
|
||
|
|
outcomes = []
|
||
|
|
for i in range(len(batch)):
|
||
|
|
if responses[i] is not None:
|
||
|
|
outcomes.append((sample,responses[i]))
|
||
|
|
|
||
|
|
with jsonlines.open(config["dataset"]["output_path"], mode='a') as writer:
|
||
|
|
for ins,result in outcomes:
|
||
|
|
gen_data = {"instruction": ins, "output": result}
|
||
|
|
writer.write(gen_data)
|