online version 1

build_openprompt.py

@@ -1,46 +1,50 @@
 import csv
-import random
+import pandas as pd
 import json
-import numpy as np
+import random
+
+from tqdm import tqdm
 
-from sklearn.model_selection import ShuffleSplit
 
 samples = {
             "x": [],
             "y": [],
         }
 little = False
-all_loaded_sample = 500000
-# 二十万条
+all_loaded_sample = 400000
+
+s_pro = all_loaded_sample / 1e+7
+# 读取概率
 with open("./data/prompts.csv") as f:
     csv_reader = csv.DictReader(f)
-    for row_number, row in enumerate(csv_reader):
-        # if row_number == random.randint(0, 1000):
-        #     break
+    process_reader = tqdm(enumerate(csv_reader))
+    for row_number, row in process_reader:
+        num_samples = len(samples['x'])
+        process_reader.set_description(f"got data num: {num_samples}")
+        if random.uniform(0, 1) > s_pro:
+            continue
         if little:
-            if row_number > 100:
+            if len(samples["x"]) > 100:
                 break
-        if row_number > all_loaded_sample:
+        if len(samples["x"]) > all_loaded_sample:
             break
         
         datum = row
+        prompt = datum['prompt']
         modifiers = json.loads(datum['raw_data'])['modifiers']
-        n = random.randint(1, 11)
-        if len(modifiers) < 3:
+        if len(modifiers) < 4:
             continue
-        label = ",".join(modifiers) if len(modifiers) > 1 else modifiers[0]
-        if 0<n and n<=6:
-            x = modifiers[0]
-        elif n>6 and n<=9:
-            x = ",".join(modifiers[:2])
-        else:
-            x = ",".join(modifiers[:3])
+        
+        # TODO: 外挂一个entity识别，过滤掉存在entity实体的数据
+        
+        label = prompt
+        x = prompt
         # 小文本到大文本，因此x更小，同时x按照6:3:1的比例分配
 
         samples["x"].append(x)
         samples["y"].append(label)
         
         
-with open("./data/dataset_openprompt.json", "w") as f:
+with open(f"./data/dataset_openprompt.json", "w") as f:
     json.dump(samples, f, indent=4, ensure_ascii=False)
-    print("*"*40, "save train done.", "with little" if little else "", "*"*40)
+    print("*"*40, f"save {num_samples} train samples done.", "with little" if little else "", "*"*40)

gpt2_generation.py

@@ -1,32 +1,10 @@
 #!/usr/bin/env python
 # coding=utf-8
-# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
-# Copyright (c) 2018, NVIDIA CORPORATION.  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.
-""" Conditional text generation with the auto-regressive models of the library (GPT/GPT-2/CTRL/Transformer-XL/XLNet)
-"""
-
-
-import argparse
 import inspect
-import time
 import logging
 from typing import Tuple
 
 import torch
-from accelerate import PartialState
-from accelerate.utils import set_seed
 
 from transformers import (
     AutoTokenizer,
@@ -49,17 +27,16 @@ from transformers import (
     XLMWithLMHeadModel,
     XLNetLMHeadModel,
     XLNetTokenizer,
+    AutoModelForSeq2SeqLM,
 )
 from transformers.modeling_outputs import CausalLMOutputWithPast
-
+from forbidden import FORBIDDEN_NOUN
 
 logging.basicConfig(
     format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
     datefmt="%m/%d/%Y %H:%M:%S",
     level=logging.INFO,
 )
-logger = logging.getLogger(__name__)
-
 MAX_LENGTH = int(10000)  # Hardcoded max length to avoid infinite loop
 
 MODEL_CLASSES = {
@@ -75,33 +52,33 @@ MODEL_CLASSES = {
     "opt": (OPTForCausalLM, GPT2Tokenizer),
 }
 
-# Padding text to help Transformer-XL and XLNet with short prompts as proposed by Aman Rusia
-# in https://github.com/rusiaaman/XLNet-gen#methodology
-# and https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e
-PREFIX = """In 1991, the remains of Russian Tsar Nicholas II and his family
-(except for Alexei and Maria) are discovered.
-The voice of Nicholas's young son, Tsarevich Alexei Nikolaevich, narrates the
-remainder of the story. 1883 Western Siberia,
-a young Grigori Rasputin is asked by his father and a group of men to perform magic.
-Rasputin has a vision and denounces one of the men as a horse thief. Although his
-father initially slaps him for making such an accusation, Rasputin watches as the
-man is chased outside and beaten. Twenty years later, Rasputin sees a vision of
-the Virgin Mary, prompting him to become a priest. Rasputin quickly becomes famous,
-with people, even a bishop, begging for his blessing. <eod> </s> <eos>"""
 
+FORBIDDEN_NOUN = set(FORBIDDEN_NOUN)
+
+class Translator:
+    def __init__(self, model_name):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+    def translate(self, text):
+        inputs = self.tokenizer(text, return_tensors="pt", padding=True)
+        outputs = self.model.generate(**inputs)
+        translated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return translated_text
+
+    def __call__(self, text):
+        return self.translate(text)
 
 #
 # Functions to prepare models' input
 #
-
-
 def prepare_ctrl_input(args, _, tokenizer, prompt_text):
-    if args.temperature > 0.7:
-        logger.info("CTRL typically works better with lower temperatures (and lower top_k).")
+    if args["temperature"] > 0.7:
+        pass
 
     encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False)
     if not any(encoded_prompt[0] == x for x in tokenizer.control_codes.values()):
-        logger.info("WARNING! You are not starting your generation from a control code so you won't get good results")
+        pass
     return prompt_text
 
 
@@ -112,8 +89,8 @@ def prepare_xlm_input(args, model, tokenizer, prompt_text):
     use_lang_emb = hasattr(model.config, "use_lang_emb") and model.config.use_lang_emb
     if hasattr(model.config, "lang2id") and use_lang_emb:
         available_languages = model.config.lang2id.keys()
-        if args.xlm_language in available_languages:
-            language = args.xlm_language
+        if args["xlm_language"] in available_languages:
+            language = args["xlm_language"]
         else:
             language = None
             while language not in available_languages:
@@ -132,13 +109,13 @@ def prepare_xlm_input(args, model, tokenizer, prompt_text):
 
 
 def prepare_xlnet_input(args, _, tokenizer, prompt_text):
-    prefix = args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
+    prefix = args["prefix"] if args["prefix"] else args["padding_text"] if args["padding_text"] else ""
     prompt_text = prefix + prompt_text
     return prompt_text
 
 
 def prepare_transfoxl_input(args, _, tokenizer, prompt_text):
-    prefix = args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
+    prefix = args["prefix"] if args["prefix"] else args["padding_text"] if args["padding_text"] else ""
     prompt_text = prefix + prompt_text
     return prompt_text
 
@@ -284,170 +261,108 @@ class _ModelFallbackWrapper(GenerationMixin):
         return self._default._reorder_cache(past_key_values, beam_idx)
 
 
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--model_type",
-        default="gpt2",
-        type=str,
-        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
-    )
-    parser.add_argument(
-        "--model_name_or_path",
-        default="./output/gpt2_openprpmpt/checkpoint-218500",
-        type=str,
-        help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
-    )
-
-    parser.add_argument("--prompt", type=str, default="")
-    parser.add_argument("--length", type=int, default=60)
-    parser.add_argument("--stop_token", type=str, default=None, help="Token at which text generation is stopped")
-
-    parser.add_argument(
-        "--temperature",
-        type=float,
-        default=1.0,
-        help="temperature of 1.0 has no effect, lower tend toward greedy sampling",
-    )
-    parser.add_argument(
-        "--repetition_penalty", type=float, default=1.0, help="primarily useful for CTRL model; in that case, use 1.2"
-    )
-    parser.add_argument("--k", type=int, default=3)
-    parser.add_argument("--p", type=float, default=0.9)
-
-    parser.add_argument("--prefix", type=str, default="", help="Text added prior to input.")
-    parser.add_argument("--padding_text", type=str, default="", help="Deprecated, the use of `--prefix` is preferred.")
-    parser.add_argument("--xlm_language", type=str, default="", help="Optional language when used with the XLM model.")
-
-    parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
-    parser.add_argument(
-        "--use_cpu",
-        action="store_true",
-        help="Whether or not to use cpu. If set to False, " "we will use gpu/npu or mps device if available",
-    )
-    parser.add_argument("--num_return_sequences", type=int, default=4, help="The number of samples to generate.")
-    parser.add_argument(
-        "--fp16",
-        action="store_true",
-        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
-    )
-    parser.add_argument("--jit", action="store_true", help="Whether or not to use jit trace to accelerate inference")
-    args = parser.parse_args()
-
-    # Initialize the distributed state.
-    distributed_state = PartialState(cpu=args.use_cpu)
-
-    logger.warning(f"device: {distributed_state.device}, 16-bits inference: {args.fp16}")
-
-    if args.seed is not None:
-        set_seed(args.seed)
-
-    # Initialize the model and tokenizer
-    try:
-        args.model_type = args.model_type.lower()
-        model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
-    except KeyError:
-        raise KeyError("the model {} you specified is not supported. You are welcome to add it and open a PR :)")
-
-    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path, padding_side='left')
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-        tokenizer.mask_token = tokenizer.eos_token
-    model = model_class.from_pretrained(args.model_name_or_path)
-
-    # Set the model to the right device
-    model.to(distributed_state.device)
-
-    if args.fp16:
-        model.half()
+def generate_prompt(
+        prompt_text,
+        args,
+        zh_en_translator,
+        nlp,
+        model,
+        tokenizer,
+        distributed_state,
+    ):
+    
     max_seq_length = getattr(model.config, "max_position_embeddings", 0)
-    args.length = adjust_length_to_model(args.length, max_sequence_length=max_seq_length)
-    logger.info(args)
-
-    prompt_text = args.prompt if args.prompt else input("Model prompt >>> ")
-
-    # Different models need different input formatting and/or extra arguments
-    requires_preprocessing = args.model_type in PREPROCESSING_FUNCTIONS.keys()
-    if requires_preprocessing:
-        prepare_input = PREPROCESSING_FUNCTIONS.get(args.model_type)
-        preprocessed_prompt_text = prepare_input(args, model, tokenizer, prompt_text)
-
-        if model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
-            tokenizer_kwargs = {"add_space_before_punct_symbol": True}
+    args["length"] = adjust_length_to_model(args["length"], max_sequence_length=max_seq_length)
+    while(1):
+        prompt_text = zh_en_translator(prompt_text)
+        # only support single input.
+        
+        # Different models need different input formatting and/or extra arguments
+        requires_preprocessing = args["model_type"] in PREPROCESSING_FUNCTIONS.keys()
+        if requires_preprocessing:
+            prepare_input = PREPROCESSING_FUNCTIONS.get(args["model_type"])
+            preprocessed_prompt_text = prepare_input(args, model, tokenizer, prompt_text)
+
+            if model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
+                tokenizer_kwargs = {"add_space_before_punct_symbol": True}
+            else:
+                tokenizer_kwargs = {}
+
+            encoded_prompt = tokenizer.encode(
+                preprocessed_prompt_text, add_special_tokens=False, return_tensors="pt", **tokenizer_kwargs
+            )
         else:
-            tokenizer_kwargs = {}
-
-        encoded_prompt = tokenizer.encode(
-            preprocessed_prompt_text, add_special_tokens=False, return_tensors="pt", **tokenizer_kwargs
-        )
-    else:
-        prefix = args.prefix if args.prefix else args.padding_text
-        encoded_prompt = tokenizer.encode(prefix + prompt_text, add_special_tokens=False, return_tensors="pt")
-    encoded_prompt = encoded_prompt.to(distributed_state.device)
+            prefix = args["prefix"] if args["prefix"] else args["padding_text"]
+            encoded_prompt = tokenizer.encode(prefix + prompt_text, add_special_tokens=False, return_tensors="pt")
+        encoded_prompt = encoded_prompt.to(distributed_state.device)
 
-    if encoded_prompt.size()[-1] == 0:
-        input_ids = None
-    else:
-        input_ids = encoded_prompt
-
-    if args.jit:
-        jit_input_texts = ["enable jit"]
-        jit_inputs = prepare_jit_inputs(jit_input_texts, model, tokenizer)
-        torch._C._jit_set_texpr_fuser_enabled(False)
-        model.config.return_dict = False
-        if hasattr(model, "forward"):
-            sig = inspect.signature(model.forward)
+        if encoded_prompt.size()[-1] == 0:
+            input_ids = None
         else:
-            sig = inspect.signature(model.__call__)
-        jit_inputs = tuple(jit_inputs[key] for key in sig.parameters if jit_inputs.get(key, None) is not None)
-        traced_model = torch.jit.trace(model, jit_inputs, strict=False)
-        traced_model = torch.jit.freeze(traced_model.eval())
-        traced_model(*jit_inputs)
-        traced_model(*jit_inputs)
-
-        model = _ModelFallbackWrapper(traced_model, model)
-    t1 = time.time()
-    output_sequences = model.generate(
-        input_ids=input_ids,
-        max_length=args.length + len(encoded_prompt[0]),
-        temperature=args.temperature,
-        top_k=args.k,
-        top_p=args.p,
-        repetition_penalty=args.repetition_penalty,
-        do_sample=True,
-        num_return_sequences=args.num_return_sequences,
-    )
-
-    # Remove the batch dimension when returning multiple sequences
-    if len(output_sequences.shape) > 2:
-        output_sequences.squeeze_()
-    
-    generated_sequences = []
-
-    for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
-        print(f"=== GENERATED SEQUENCE {generated_sequence_idx + 1} ===")
-        generated_sequence = generated_sequence.tolist()
-
-        # Decode text
-        text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
-
-        # Remove all text after the stop token
-        text = text[: text.find(args.stop_token) if args.stop_token else None]
-
-        # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
-        total_sequence = (
-            prompt_text + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
-        )
-
-        generated_sequences.append(total_sequence)
-        print(total_sequence)
-
-    t2 = time.time()
-    print("*"*60)
-    print(f"Time cost: {t2-t1}")
-
-    return generated_sequences
+            input_ids = encoded_prompt
+
+        if args["jit"]:
+            jit_input_texts = ["enable jit"]
+            jit_inputs = prepare_jit_inputs(jit_input_texts, model, tokenizer)
+            torch._C._jit_set_texpr_fuser_enabled(False)
+            model.config.return_dict = False
+            if hasattr(model, "forward"):
+                sig = inspect.signature(model.forward)
+            else:
+                sig = inspect.signature(model.__call__)
+            jit_inputs = tuple(jit_inputs[key] for key in sig.parameters if jit_inputs.get(key, None) is not None)
+            traced_model = torch.jit.trace(model, jit_inputs, strict=False)
+            traced_model = torch.jit.freeze(traced_model.eval())
+            traced_model(*jit_inputs)
+            traced_model(*jit_inputs)
+
+            model = _ModelFallbackWrapper(traced_model, model)
+            
+        generated_sequences = []
+
+        for generated_sequence_idx in range(args["num_return_sequences"]):
+            repeat_gen_time = 0
+            while(1):
+                repeat_gen_time = repeat_gen_time + 1
+                generated_sequence = model.generate(
+                    input_ids=input_ids,
+                    max_length=args["length"] + len(encoded_prompt[0]),
+                    temperature=args["temperature"],
+                    top_k=args["k"],
+                    top_p=args["p"],
+                    repetition_penalty=args["repetition_penalty"],
+                    do_sample=True,
+                    num_return_sequences=1,
+                    pad_token_id=tokenizer.pad_token_id
+                )
+                # Remove the n_sequence dimension when returning single sequence
+                if len(generated_sequence.shape) >1:
+                    generated_sequence.squeeze_()
+                
+                generated_sequence = generated_sequence.tolist()
+
+                # Decode text
+                text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
+
+                # Remove all text after the stop token
+                text = text[: text.find(args["stop_token"]) if args["stop_token"] else None]
+
+                # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
+                total_sequence = (
+                    prompt_text + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
+                )
+                # no checking for prompt_text.
+                docs = nlp(text)
+                nouns = [token.text for token in docs if token.pos_ == 'NOUN']
+                nouns = set(nouns)
+                if nouns.intersection(FORBIDDEN_NOUN) and repeat_gen_time < 10:
+                    continue
+                else:
+                    break
+            generated_sequences.append(total_sequence)
+
+        return generated_sequences
 
 
 if __name__ == "__main__":
-    main()
+    generate_prompt()

sft.py

@@ -1,10 +1,12 @@
 import time
 import evaluate
 import numpy as np
-
+import math
 from transformers import DataCollatorForLanguageModeling, DataCollatorForSeq2Seq
 from transformers import TrainingArguments, Trainer
 
+from transformers.trainer_callback import TrainerCallback
+
 from utils import (
     get_dataset,
     get_tok_and_model,
@@ -17,34 +19,27 @@ tokenizer, model = get_tok_and_model(f"./models/{base_model}")
 tokenizer.pad_token = tokenizer.eos_token
 rouge = evaluate.load("rouge")
 
-# train_data, test_data = get_open_prompt_data("./data")
-# train_dataset, test_dataset = get_dataset(train_data, test_data)
 dict_data = get_dict_dataset("./data")
 dataset = get_advance_dataset(dict_data)
-dataset = dataset.train_test_split(test_size=0.2)
+dataset = dataset.train_test_split(test_size=0.05)
 
 def preprocess_function(examples):
     x_inputs = [x for x in examples["x"]]
     y_inputs = examples["y"]
-    model_inputs = tokenizer(x_inputs, max_length=128, truncation=True)
+    model_inputs = tokenizer(x_inputs, max_length=256, truncation=True)
 
-    labels = tokenizer(text_target=y_inputs, max_length=128, truncation=True)
+    labels = tokenizer(y_inputs, max_length=256, truncation=True)
 
-    model_inputs["labels"] = model_inputs["input_ids"]
+    model_inputs["labels"] = labels["input_ids"]
     return model_inputs
 
-def compute_metrics(eval_pred):
-    predictions, labels = eval_pred
-    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
-    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
-    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
-
-    result = rouge.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
+class CustomCallback(TrainerCallback):
+    def on_epoch_end(self, args, state, control, **kwargs):
+        control.should_evaluate=True,
 
-    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]
-    result["gen_len"] = np.mean(prediction_lens)
-
-    return {k: round(v, 4) for k, v in result.items()}
+    def on_evaluate(self, args, state, control, **kwargs):
+        eval_results = kwargs["metrics"]
+        print(f"Perplexity: {math.exp(eval_results['eval_loss']):.2f}\n")
 
 # data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
@@ -69,7 +64,7 @@ training_args = TrainingArguments(
     adam_beta1=0.9,
     adam_beta2=0.98,
     save_total_limit=1,
-    num_train_epochs=100,
+    num_train_epochs=80,
     fp16=True,
     push_to_hub=False,
 )
@@ -81,12 +76,8 @@ trainer = Trainer(
     eval_dataset=tokenized_dataset["test"],
     tokenizer=tokenizer,
     data_collator=data_collator,
+    callbacks=[CustomCallback]
 )
 
 trainer.train()
 
-import math
-
-eval_results = trainer.evaluate()
-print(f"Perplexity: {math.exp(eval_results['eval_loss']):.2f}")
-