init commit

easydistill/rank/train.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
+from jinja2 import Environment, BaseLoader, FileSystemLoader
+from datasets import load_dataset, Dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import DPOTrainer, DPOConfig
+import copy
+
+
+def process_dataset(dataset_path, dataset_seed, env, template):
+    examples = []
+    with open(dataset_path, 'r') as file:
+        examples = json.load(file)
+    output_text = {
+        "prompt": [],
+        "chosen": [],
+        "rejected": []
+    }
+    # use chat template
+    for i in range(len(examples)):
+        try:
+            prompt_message = {"content": examples[i]["prompt"]}
+            prompt = template.render(message=prompt_message, add_generation_prompt=False, add_output=False)
+
+            chosen_message = {"content": examples[i]["prompt"], "output": examples[i]["chosen"]}
+            chosen = template.render(message=chosen_message, add_generation_prompt=False, add_output=True)
+            chosen = chosen[len(prompt):]
+
+            rejected_message = {"content": examples[i]["prompt"], "output": examples[i]["rejected"]}
+            rejected = template.render(message=rejected_message, add_generation_prompt=False, add_output=True)
+            rejected = rejected[len(prompt):]
+
+            output_text["prompt"].append(prompt)
+            output_text["chosen"].append(chosen)
+            output_text["rejected"].append(rejected)
+        except:
+            logging.warning(f"Error processing sample.")
+            
+    dataset = Dataset.from_dict(output_text)
+    dataset = dataset.shuffle(seed=dataset_seed)        
+    return dataset
+
+
+def train(config):    
+    dataset_path = config["dataset"]["labeled_path"]
+    dataset_seed = config["dataset"]["seed"]
+    
+    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)
+    dataset = process_dataset(dataset_path, dataset_seed, env, template)
+    
+    student_tokenizer = AutoTokenizer.from_pretrained(
+        config["models"]["student"], 
+        trust_remote_code=True
+    )
+    student_model = AutoModelForCausalLM.from_pretrained(
+        config["models"]["student"],
+        trust_remote_code=True
+    )    
+    
+    training_arguments = DPOConfig(**config["training"])    
+    trainer = DPOTrainer(
+        student_model,
+        ref_model=copy.deepcopy(student_model),
+        args=training_arguments,
+        train_dataset=dataset,
+        processing_class=student_tokenizer
+    )
+
+    trainer.train()
+    trainer.save_model(config["training"]["output_dir"])
+    student_tokenizer.save_pretrained(config["training"]["output_dir"])
+
+
+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))
+    train(config)
+
+
+if __name__ == "__main__":
+    main()