init commit

easydistill/synthesis/cot_synthesis.py

@@ -0,0 +1,274 @@
+
+# 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)