init commit

recipes/distilqwen_series/distillqwen2/dpo_student_infer_only.py

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