Delete main.py with huggingface_hub

main.py

@@ -1,203 +0,0 @@
-
-import os
-import time
-import wandb
-import torch
-import argparse
-from datasets import load_dataset
-from typing import List, Dict, Union
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    TrainingArguments,
-    DataCollatorForLanguageModeling
-)
-
-from src.args import default_args
-from src.orpo_trainer import ORPOTrainer
-from src.utils import preprocess_logits_for_metrics, dataset_split_selector
-
-class ORPO(object):
-    def __init__(self, args) -> None:
-        self.start = time.gmtime()
-        self.args = args
-
-        # Load Tokenizer
-        print(">>> 1. Loading Tokenizer")
-        self.tokenizer = AutoTokenizer.from_pretrained(self.args.model_name, cache_dir=self.args.cache_dir)
-        if self.tokenizer.chat_template is None:
-            self.tokenizer.chat_template = "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<|system|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"
-            print("     1-1. Chat Template Applied (<|user|> <|assistant|>)")
-        else:
-            pass
-        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
-
-        # Load Model
-        print(">>> 2. Loading Model")
-        if self.args.flash_attention_2:
-            self.model = AutoModelForCausalLM.from_pretrained(self.args.model_name, 
-                                                              cache_dir=self.args.cache_dir,
-                                                              torch_dtype=torch.bfloat16,
-                                                              attn_implementation="flash_attention_2")
-        else:
-            self.model = AutoModelForCausalLM.from_pretrained(self.args.model_name, 
-                                                              cache_dir=self.args.cache_dir,
-                                                              torch_dtype=torch.bfloat16)
-                                                          
-        # Load Dataset
-        print(">>> 3. Loading Dataset")
-        self.data = load_dataset(self.args.data_name, cache_dir=self.args.cache_dir)
-
-        # Preprocess Dataset
-        print(">>> 4. Filtering and Preprocessing Dataset")
-        data_split = dataset_split_selector(self.data)
-
-        if len(data_split) == 1:
-            self.is_test = False
-            train_split = data_split[0]
-            print(f"   >>> Test Set = {self.is_test}")
-        else:
-            self.is_test = True
-            train_split = data_split[0]
-            test_split = data_split[1]
-
-            test = self.data[test_split].filter(self.filter_dataset)
-            self.test = test.map(self.preprocess_dataset, batched=True, num_proc=self.args.num_proc, remove_columns=self.data[test_split].column_names)       
-
-        train = self.data[train_split].filter(self.filter_dataset)
-        print(f"\n\n>>> {len(train)} / {len(self.data[train_split])} rows left after filtering by prompt length.")
-        self.train = train.map(self.preprocess_dataset, batched=True, num_proc=self.args.num_proc, remove_columns=self.data[train_split].column_names)                       
-                
-        # Set WANDB & Logging Configurations
-        self.run_name = f"{self.args.model_name.split('/')[-1]}-{self.args.data_name.split('/')[-1]}-lambda{self.args.alpha}-ORPO-{self.start.tm_mday}-{self.start.tm_hour}-{self.start.tm_min}"
-        self.save_dir = os.path.join('./checkpoints/', f"{self.args.data_name.split('/')[-1]}/{self.run_name}")
-        self.log_dir = os.path.join('./checkpoints/', f"{self.args.data_name.split('/')[-1]}/{self.run_name}/logs")
-        
-        os.makedirs(self.save_dir, exist_ok=True)
-        os.makedirs(self.log_dir, exist_ok=True)
-
-    def preprocess_dataset(self, examples: Union[List, Dict]):
-        if ('instruction' in examples.keys()) or ('question' in examples.keys()):
-            prompt_key = 'instruction' if 'instruction' in examples.keys() else 'question'
-            prompt = [self.tokenizer.apply_chat_template([{'role': 'user', 'content': item}], tokenize=False, add_generation_prompt=True) for item in examples[prompt_key]]
-            chosen = [self.tokenizer.apply_chat_template([{'role': 'user', 'content': item_prompt}, {'role': 'assistant', 'content': item_chosen}], tokenize=False) for item_prompt, item_chosen in zip(examples[prompt_key], examples['chosen'])]
-            rejected = [self.tokenizer.apply_chat_template([{'role': 'user', 'content': item_prompt}, {'role': 'assistant', 'content': item_rejected}], tokenize=False) for item_prompt, item_rejected in zip(examples[prompt_key], examples['rejected'])]
-        else:
-            prompt = [self.tokenizer.apply_chat_template([item[0]], tokenize=False, add_generation_prompt=True) for item in examples['chosen']]
-            chosen = [self.tokenizer.apply_chat_template(item, tokenize=False) for item in examples['chosen']]
-            rejected = [self.tokenizer.apply_chat_template(item, tokenize=False) for item in examples['rejected']]
-    
-        model_inputs = self.tokenizer(prompt,
-                                      max_length=self.args.response_max_length,
-                                      padding='max_length',
-                                      truncation=True,
-                                      return_tensors='pt')
-        pos_labels = self.tokenizer(chosen,
-                                    max_length=self.args.response_max_length,
-                                    padding='max_length',
-                                    truncation=True,
-                                    return_tensors='pt')
-        neg_labels = self.tokenizer(rejected,
-                                    max_length=self.args.response_max_length,
-                                    padding='max_length',
-                                    truncation=True,
-                                    return_tensors='pt') 
-                
-        model_inputs['positive_input_ids'] = pos_labels['input_ids']
-        model_inputs['positive_attention_mask'] = pos_labels['attention_mask']
-        
-        model_inputs['negative_input_ids'] = neg_labels['input_ids']
-        model_inputs['negative_attention_mask'] = neg_labels['attention_mask']
-        
-        return model_inputs
-
-    def filter_dataset(self, examples: Union[List, Dict]):
-        if 'instruction' in examples.keys():
-            query = examples['instruction']
-            prompt_length = self.tokenizer.apply_chat_template([{'content': query, 'role': 'user'}], tokenize=True, add_generation_prompt=True, return_tensors='pt').size(-1)
-        elif 'question' in examples.keys():
-            query = examples['question']
-            prompt_length = self.tokenizer.apply_chat_template([{'content': query, 'role': 'user'}], tokenize=True, add_generation_prompt=True, return_tensors='pt').size(-1)
-        else:
-            prompt_length = self.tokenizer.apply_chat_template([examples['chosen'][0]], tokenize=True, add_generation_prompt=True, return_tensors='pt').size(-1)  
-           
-        if prompt_length < self.args.prompt_max_length:    
-            return True
-        else:
-            return False
-
-    def prepare_trainer(self):
-        wandb.init(name=self.run_name)
-        arguments = TrainingArguments(
-            output_dir=self.save_dir,  # The output directory
-            logging_dir=self.log_dir,
-            logging_steps=50,
-            learning_rate=self.args.lr,
-            overwrite_output_dir=True,  # overwrite the content of the output directory
-            num_train_epochs=self.args.num_train_epochs,  # number of training epochs
-            per_device_train_batch_size=self.args.per_device_train_batch_size,  # batch size for training
-            per_device_eval_batch_size=self.args.per_device_eval_batch_size,  # batch size for evaluation
-            evaluation_strategy=self.args.evaluation_strategy if self.is_test else 'no',  # batch size for evaluation
-            save_strategy=self.args.evaluation_strategy,
-            optim=self.args.optim,
-            warmup_steps=self.args.warmup_steps,
-            gradient_accumulation_steps=self.args.gradient_accumulation_steps,
-            gradient_checkpointing=True, #if ('llama' in self.args.model_name.lower()) or ('mistral' in self.args.model_name.lower()) else False,
-            gradient_checkpointing_kwargs={'use_reentrant':True},
-            load_best_model_at_end=self.is_test,
-            do_train=True,
-            do_eval=self.is_test,
-            lr_scheduler_type=self.args.lr_scheduler_type,
-            remove_unused_columns=False,
-            report_to='wandb',
-            run_name=self.run_name,
-            bf16=True
-        )
-        
-        data_collator = DataCollatorForLanguageModeling(tokenizer=self.tokenizer, mlm=False)
-        
-        self.trainer = ORPOTrainer(
-            model=self.model,
-            alpha=self.args.alpha,
-            pad=self.tokenizer.pad_token_id,
-            args=arguments,
-            train_dataset=self.train,
-            eval_dataset=self.test if self.is_test else None,
-            data_collator=data_collator,
-            preprocess_logits_for_metrics=preprocess_logits_for_metrics
-        )
-        
-    def run(self):
-        print(">>> 5. Preparing ORPOTrainer")
-        self.prepare_trainer()
-        self.trainer.train()
-
-        # Saving code for FSDP
-        if self.trainer.is_fsdp_enabled:
-            self.trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
-        self.trainer.save_model()
-        
-        
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser("ORPO")
-    args = default_args(parser)
-
-    # Set WANDB configurations
-    if args.wandb_entity is not None and args.wandb_project_name is not None:
-        os.environ["WANDB_ENTITY"] = args.wandb_entity
-        os.environ["WANDB_PROJECT"] = args.wandb_project_name
-    else:
-        pass
-    os.environ["TOKENIZERS_PARALLELISM"] = 'false'
-
-    print("================================================================================================\n")
-    print(f">>> Fine-tuning {args.model_name} with ORPO on {args.data_name}\n")
-    print("================================================================================================")
-    print("\n\n>>> Summary:")
-    print(f"    - Lambda              : {args.alpha}")
-    print(f"    - Training Epochs     : {args.num_train_epochs}")
-    print(f"    - Prompt Max Length   : {args.prompt_max_length}")
-    print(f"    - Response Max Length : {args.response_max_length}")
-
-    item = ORPO(args=args)
-    item.run()