Delete train.py

train.py

@@ -1,385 +0,0 @@
-import json
-import torch
-import random
-import transformers
-import networkx as nx
-from tqdm import tqdm
-from peft import (LoraConfig, get_peft_model,
-                  prepare_model_for_kbit_training)
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-
-
-
-class QloraTrainer_CS:
-    def __init__(self, config: dict, use_predefined_graph=False):
-        self.config = config
-        self.use_predefined_graph = use_predefined_graph
-        self.tokenizer = None
-        self.base_model = None
-        self.adapter_model = None
-        self.merged_model = None
-        self.transformer_trainer = None
-        self.test_data = None
-
-        template_file_path = 'configs/alpaca.json'
-        with open(template_file_path) as fp:
-            self.template = json.load(fp)
-
-
-    def load_base_model(self):
-        model_id = self.config['inference']["base_model"]
-        print(model_id)
-
-        bnb_config = BitsAndBytesConfig(
-            load_in_8bit=True,
-            bnb_8bit_use_double_quant=True,
-            bnb_8bit_quant_type="nf8",
-            bnb_8bit_compute_dtype=torch.bfloat16
-        )
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        tokenizer.model_max_length = self.config['training']['tokenizer']["max_length"]
-        if not tokenizer.pad_token:
-            tokenizer.pad_token = tokenizer.eos_token
-        model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, torch_dtype=torch.bfloat16)
-        if model.device.type != 'cuda':
-            model.to('cuda')
-
-        model.gradient_checkpointing_enable()
-        model = prepare_model_for_kbit_training(model)
-
-        self.tokenizer = tokenizer
-        self.base_model = model
-
-
-    def train(self):        
-        # Set up lora config or load pre-trained adapter
-        lora_config = LoraConfig(
-            r=self.config['training']['qlora']['rank'],
-            lora_alpha=self.config['training']['qlora']['lora_alpha'],
-            target_modules=self.config['training']['qlora']['target_modules'],
-            lora_dropout=self.config['training']['qlora']['lora_dropout'],
-            bias="none",
-            task_type="CAUSAL_LM",
-        )
-        model = get_peft_model(self.base_model, lora_config)
-        self._print_trainable_parameters(model)
-
-        print("Start data preprocessing")
-        train_data = self._process_data_instruction()
-
-        print('Length of dataset: ', len(train_data))
-
-        print("Start training")
-        self.transformer_trainer = transformers.Trainer(
-            model=model,
-            train_dataset=train_data,
-            args=transformers.TrainingArguments(
-                per_device_train_batch_size=self.config["training"]['trainer_args']["per_device_train_batch_size"],
-                gradient_accumulation_steps=self.config['model_saving']['index'],
-                warmup_steps=self.config["training"]['trainer_args']["warmup_steps"],
-                num_train_epochs=self.config["training"]['trainer_args']["num_train_epochs"],
-                learning_rate=self.config["training"]['trainer_args']["learning_rate"],
-                lr_scheduler_type=self.config["training"]['trainer_args']["lr_scheduler_type"],
-                fp16=self.config["training"]['trainer_args']["fp16"],
-                logging_steps=self.config["training"]['trainer_args']["logging_steps"],
-                output_dir=self.config["training"]['trainer_args']["trainer_output_dir"],
-                report_to="wandb",
-                save_steps=self.config["training"]['trainer_args']["save_steps"],
-            ),
-            data_collator=transformers.DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
-        )
-
-        model.config.use_cache = False
-
-        self.transformer_trainer.train()
-
-        model_save_path = f"{self.config['model_saving']['model_output_dir']}/{self.config['model_saving']['model_name']}_{self.config['model_saving']['index']}_adapter_test_graph"
-        self.transformer_trainer.save_model(model_save_path)
-
-        self.adapter_model = model
-        print(f"Training complete, adapter model saved in {model_save_path}")
-
-
-    def _print_trainable_parameters(self, model):
-        """
-        Prints the number of trainable parameters in the model.
-        """
-        trainable_params = 0
-        all_param = 0
-        for _, param in model.named_parameters():
-            all_param += param.numel()
-            if param.requires_grad:
-                trainable_params += param.numel()
-        print(
-            f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
-        )
-
-
-    def _process_data_instruction(self):
-        context_window = self.tokenizer.model_max_length
-        if self.use_predefined_graph:
-            graph_data = nx.read_gexf('datasets/' + self.config["training"]["predefined_graph_path"], node_type=None, relabel=False, version='1.2draft')
-        else:
-            graph_path = self.config['data_downloading']['download_directory'] + 'description/' + self.config['data_downloading']['gexf_file']
-            graph_data = nx.read_gexf(graph_path, node_type=None, relabel=False, version='1.2draft')
-        raw_graph = graph_data
-
-        test_set_size = len(graph_data.nodes()) // 10
-
-        all_test_nodes = set(list(graph_data.nodes())[:test_set_size])
-        all_train_nodes = set(list(graph_data.nodes())[test_set_size:])
-
-        raw_id_2_title_abs = dict()
-        for paper_id in list(graph_data.nodes())[test_set_size:]:
-            title = graph_data.nodes()[paper_id]['title']
-            abstract = graph_data.nodes()[paper_id]['abstract']
-            raw_id_2_title_abs[paper_id] = [title, abstract]
-
-        raw_id_2_intro = dict()
-        for paper_id in list(graph_data.nodes())[test_set_size:]:
-            if graph_data.nodes[paper_id]['introduction'] != '':
-                intro = graph_data.nodes[paper_id]['introduction']
-                raw_id_2_intro[paper_id] = intro
-
-        raw_id_pair_2_sentence = dict()
-        for edge in list(graph_data.edges()):
-            sentence = graph_data.edges()[edge]['sentence']
-            raw_id_pair_2_sentence[edge] = sentence
-
-
-        test_data = []
-        edge_list = []
-        for edge in list(raw_graph.edges()):
-            src, tar = edge
-            if src not in all_test_nodes and tar not in all_test_nodes:
-                edge_list.append(edge)
-            else:
-                test_data.append(edge)
-        train_num = int(len(edge_list))
-
-        data_LP = []
-        data_abstract_2_title = []
-        data_paper_retrieval = []
-        data_citation_sentence = []
-        data_abs_completion = []
-        data_title_2_abs = []
-        data_intro_2_abs = []
-
-
-        for sample in tqdm(random.sample(edge_list, train_num)):
-            source, target = sample[0], sample[1]
-            source_title, source_abs = raw_id_2_title_abs[source]
-            target_title, target_abs = raw_id_2_title_abs[target]
-            # LP prompt
-            rand_ind = random.choice(list(raw_id_2_title_abs.keys()))
-            neg_title, neg_abs = raw_id_2_title_abs[rand_ind]
-            data_LP.append({'s_title':source_title, 's_abs':source_abs, 't_title':target_title, 't_abs':target_abs, 'label':'yes'})
-            data_LP.append({'s_title':source_title, 's_abs':source_abs, 't_title':neg_title, 't_abs':neg_abs, 'label':'no'})
-        
-        for sample in tqdm(random.sample(edge_list, train_num)):
-            source, target = sample[0], sample[1]
-            source_title, source_abs = raw_id_2_title_abs[source]
-            target_title, target_abs = raw_id_2_title_abs[target]
-            # abs_2_title prompt
-            data_abstract_2_title.append({'title':source_title, 'abs':source_abs})
-            data_abstract_2_title.append({'title':target_title, 'abs':target_abs})
-
-        for sample in tqdm(random.sample(edge_list, train_num)):
-            source, target = sample[0], sample[1]
-            source_title, source_abs = raw_id_2_title_abs[source]
-            target_title, target_abs = raw_id_2_title_abs[target]
-            # paper_retrieval prompt
-            neighbors = list(nx.all_neighbors(raw_graph, source))
-            sample_node_list = list(all_train_nodes - set(neighbors) - set([source]) - set([target]))
-            sampled_neg_nodes = random.sample(sample_node_list, 5) + [target]
-            random.shuffle(sampled_neg_nodes)
-            data_paper_retrieval.append({'title':source_title, 'abs':source_abs, 'sample_title': [raw_id_2_title_abs[node][0] for node in sampled_neg_nodes], 'right_title':target_title})
-        
-        for sample in tqdm(random.sample(edge_list, train_num)):
-            source, target = sample[0], sample[1]
-            source_title, source_abs = raw_id_2_title_abs[source]
-            target_title, target_abs = raw_id_2_title_abs[target]    
-            # citation_sentence prompt
-            citation_sentence = raw_id_pair_2_sentence[(source, target)] if (source, target) in raw_id_pair_2_sentence.keys() else raw_id_pair_2_sentence[(target, source)]
-            data_citation_sentence.append({'s_title':source_title, 's_abs':source_abs, 't_title':target_title, 't_abs':target_abs, 'sentence': citation_sentence})
-        
-        for sample in tqdm(random.sample(edge_list, train_num)):
-            source, target = sample[0], sample[1]
-            source_title, source_abs = raw_id_2_title_abs[source]
-            target_title, target_abs = raw_id_2_title_abs[target]
-            # abs_complete prompt
-            data_abs_completion.append({'title':source_title, 'abs':source_abs})
-            data_abs_completion.append({'title':target_title, 'abs':target_abs})
-            
-        for sample in tqdm(random.sample(edge_list, train_num)):
-            source, target = sample[0], sample[1]
-            source_title, source_abs = raw_id_2_title_abs[source]
-            target_title, target_abs = raw_id_2_title_abs[target]
-            # title_2_abs prompt
-            data_title_2_abs.append({'title':source_title, 'right_abs':source_abs})
-            data_title_2_abs.append({'title':target_title, 'right_abs':target_abs})
-
-        for sample in tqdm(random.sample(edge_list, train_num)):
-            source, target = sample[0], sample[1]
-            if source in raw_id_2_intro:
-                source_intro = raw_id_2_intro[source]
-                _, source_abs = raw_id_2_title_abs[source]
-                data_intro_2_abs.append({'intro':source_intro, 'abs':source_abs})
-            if target in raw_id_2_intro:
-                target_intro = raw_id_2_intro[target]
-                _, target_abs = raw_id_2_title_abs[target]
-                data_intro_2_abs.append({'intro':target_intro, 'abs':target_abs})            
-
-        data_prompt = []
-        data_prompt += [self._generate_paper_retrieval_prompt(data_point) for data_point in data_paper_retrieval]
-        data_prompt += [self._generate_LP_prompt(data_point) for data_point in data_LP]
-        data_prompt += [self._generate_abstract_2_title_prompt(data_point) for data_point in data_abstract_2_title]
-        data_prompt += [self._generate_citation_sentence_prompt(data_point) for data_point in data_citation_sentence]
-        data_prompt += [self._generate_abstract_completion_prompt(data_point) for data_point in data_abs_completion]
-        data_prompt += [self._generate_title_2_abstract_prompt(data_point) for data_point in data_title_2_abs]
-        data_prompt += [self._generate_intro_2_abstract_prompt(data_point, context_window) for data_point in data_intro_2_abs]
-
-        print("Total prompts:", len(data_prompt))
-        random.shuffle(data_prompt)
-        if self.tokenizer.chat_template is None:
-            data_tokenized = [self.tokenizer(sample,  max_length=context_window, truncation=True) for sample in tqdm(data_prompt)]
-        else:
-            data_tokenized = [self.tokenizer.apply_chat_template(sample,  max_length=context_window, truncation=True, tokenize=False) for sample in tqdm(data_prompt)]
-
-        return data_tokenized
-
-    
-    def _generate_LP_prompt(self, data_point: dict):
-        instruction = "Determine if paper A will cite paper B."
-
-        prompt_input = ""
-        prompt_input = prompt_input + "Title of Paper A: " + (data_point['s_title'] if data_point['s_title'] != None else 'Unknown') + "\n"
-        prompt_input = prompt_input + "Abstract of Paper A: " + (data_point['s_abs'] if data_point['s_abs'] != None else 'Unknown') + "\n"
-        prompt_input = prompt_input + "Title of Paper B: " + (data_point['t_title'] if data_point['t_title'] != None else 'Unknown') + "\n"
-        prompt_input = prompt_input + "Abstract of Paper B: " + (data_point['t_abs'] if data_point['t_abs'] != None else 'Unknown') + "\n"
-
-        if self.tokenizer.chat_template is None:
-            res = self.template["prompt_input"].format(instruction=instruction, input=prompt_input)
-            res = f"{res}{data_point['label']}"
-        else:
-            res = [
-                {"role": "user", "content": self.template["prompt_input"].format(instruction=instruction, input=prompt_input)},
-                {"role": "assistant", "content": data_point['label']}
-            ]
-
-        return res
- 
-    def _generate_abstract_2_title_prompt(self, data_point: dict):
-        instruction = "Please generate the title of paper based on its abstract."
-
-        prompt_input = ""
-        prompt_input = prompt_input + "Abstract: " + data_point['abs'] + "\n"
-
-        if self.tokenizer.chat_template is None:
-            res = self.template["prompt_input"].format(instruction=instruction, input=prompt_input)
-            res = f"{res}{data_point['title']}"
-        else:
-            res = [
-                {"role": "user", "content": self.template["prompt_input"].format(instruction=instruction, input=prompt_input)},
-                {"role": "assistant", "content": data_point['title']}
-            ]
-
-        return res
-    
-    def _generate_paper_retrieval_prompt(self, data_point: dict):
-        instruction = "Please select the paper that is more likely to be cited by paper A from candidate papers."
-        
-        prompt_input = ""
-        prompt_input = prompt_input + "Title of the Paper A: " + data_point['title'] + "\n"
-        prompt_input = prompt_input + "Abstract of the Paper A: " + data_point['abs'] + "\n"
-        prompt_input = prompt_input + "candidate papers: " + "\n"
-        for i in range(len(data_point['sample_title'])):
-            prompt_input = prompt_input + str(i) + '. ' + data_point['sample_title'][i] + "\n"
-        
-        if self.tokenizer.chat_template is None:
-            res = self.template["prompt_input"].format(instruction=instruction, input=prompt_input)
-            res = f"{res}{data_point['right_title']}"
-        else:
-            res = [
-                {"role": "user", "content": self.template["prompt_input"].format(instruction=instruction, input=prompt_input)},
-                {"role": "assistant", "content": data_point['right_title']}
-            ]
-
-        return res
-
-    def _generate_citation_sentence_prompt(self, data_point: dict):
-        instruction = "Please generate the citation sentence of how Paper A cites paper B in its related work section."
-        
-        prompt_input = ""
-        prompt_input = prompt_input + "Title of Paper A: " + (data_point['s_title'] if data_point['s_title'] != None else 'Unknown') + "\n"
-        prompt_input = prompt_input + "Abstract of Paper A: " + (data_point['s_abs'] if data_point['s_abs'] != None else 'Unknown') + "\n"
-        prompt_input = prompt_input + "Title of Paper B: " + (data_point['t_title'] if data_point['t_title'] != None else 'Unknown') + "\n"
-        prompt_input = prompt_input + "Abstract of Paper B: " + (data_point['t_abs'] if data_point['t_abs'] != None else 'Unknown') + "\n"
-
-        if self.tokenizer.chat_template is None:
-            res = self.template["prompt_input"].format(instruction=instruction, input=prompt_input)
-            res = f"{res}{data_point['sentence']}"
-        else:
-            res = [
-                {"role": "user", "content": self.template["prompt_input"].format(instruction=instruction, input=prompt_input)},
-                {"role": "assistant", "content": data_point['sentence']}
-            ]
-
-        return res
-    
-    def _generate_abstract_completion_prompt(self, data_point: dict):
-        instruction = "Please complete the abstract of a paper."
-
-        prompt_input = ""
-        prompt_input = prompt_input + "Title: " + data_point['title'] if data_point['title'] != None else 'Unknown' + "\n"
-        
-        split_abs = data_point['abs'][: int(0.3*len(data_point['abs']))]
-        prompt_input = prompt_input + "Part of abstract: " + split_abs + "\n"
-
-        if self.tokenizer.chat_template is None:
-            res = self.template["prompt_input"].format(instruction=instruction, input=prompt_input)
-            res = f"{res}{data_point['abs']}"
-        else:
-            res = [
-                {"role": "user", "content": self.template["prompt_input"].format(instruction=instruction, input=prompt_input)},
-                {"role": "assistant", "content": data_point['abs']}
-            ]
-
-        return res
-    
-    def _generate_title_2_abstract_prompt(self, data_point: dict):
-        instruction = "Please generate the abstract of paper based on its title."
-
-        prompt_input = ""
-        prompt_input = prompt_input + "Title: " + data_point['title'] + "\n"
-        
-        if self.tokenizer.chat_template is None:
-            res = self.template["prompt_input"].format(instruction=instruction, input=prompt_input)
-            res = f"{res}{data_point['right_abs']}"
-        else:
-            res = [
-                {"role": "user", "content": self.template["prompt_input"].format(instruction=instruction, input=prompt_input)},
-                {"role": "assistant", "content": data_point['right_abs']}
-            ]
-
-        return res
-    
-    def _generate_intro_2_abstract_prompt(self, data_point: dict, context_window):
-        instruction = "Please generate the abstract of paper based on its introduction section."
-
-        prompt_input = ""
-        prompt_input = prompt_input + "Introduction: " + data_point['intro'] + "\n"
-
-        # Reduce it to make it fit
-        prompt_input = prompt_input[:int(context_window*2)]
-        
-        if self.tokenizer.chat_template is None:
-            res = self.template["prompt_input"].format(instruction=instruction, input=prompt_input)
-            res = f"{res}{data_point['abs']}"
-        else:
-            res = [
-                {"role": "user", "content": self.template["prompt_input"].format(instruction=instruction, input=prompt_input)},
-                {"role": "assistant", "content": data_point['abs']}
-            ]
-
-        return res