<gh_stars>1-10
#!/usr/bin/env python3
# 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 logging

import numpy as np
import torch

from transformers import (
    CTRLLMHeadModel,
    CTRLTokenizer,
    GPT2LMHeadModel,
    GPT2Tokenizer,
    OpenAIGPTLMHeadModel,
    OpenAIGPTTokenizer,
    TransfoXLLMHeadModel,
    TransfoXLTokenizer,
    XLMTokenizer,
    XLMWithLMHeadModel,
    XLNetLMHeadModel,
    XLNetTokenizer,
)
from thop import profile
from torchstat import stat
from transformers import BartForConditionalGeneration, BartTokenizer
from transformers import T5ForConditionalGeneration, T5Tokenizer



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 = {
    "gpt2": (GPT2LMHeadModel, GPT2Tokenizer),
    "ctrl": (CTRLLMHeadModel, CTRLTokenizer),
    "openai-gpt": (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
    "xlnet": (XLNetLMHeadModel, XLNetTokenizer),
    "transfo-xl": (TransfoXLLMHeadModel, TransfoXLTokenizer),
    "xlm": (XLMWithLMHeadModel, XLMTokenizer),
}

# Padding text to help Transformer-XL and XLNet with short prompts as proposed by <NAME>
# in https://github.com/rusiaaman/XLNet-gen#methodology
# and https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e
PADDING_TEXT = """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 <NAME>, narrates the
remainder of the story. 1883 Western Siberia,
a young <NAME> 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>"""


def set_seed(args):
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


#
# 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).")

    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")
    return prompt_text


def prepare_xlm_input(args, model, tokenizer, prompt_text):
    # kwargs = {"language": None, "mask_token_id": None}

    # Set the language
    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
        else:
            language = None
            while language not in available_languages:
                language = input("Using XLM. Select language in " + str(list(available_languages)) + " >>> ")

        model.config.lang_id = model.config.lang2id[language]
        # kwargs["language"] = tokenizer.lang2id[language]

    # TODO fix mask_token_id setup when configurations will be synchronized between models and tokenizers
    # XLM masked-language modeling (MLM) models need masked token
    # is_xlm_mlm = "mlm" in args.model_name_or_path
    # if is_xlm_mlm:
    #     kwargs["mask_token_id"] = tokenizer.mask_token_id

    return prompt_text


def prepare_xlnet_input(args, _, tokenizer, prompt_text):
    prompt_text = (args.padding_text if args.padding_text else PADDING_TEXT) + prompt_text
    return prompt_text


def prepare_transfoxl_input(args, _, tokenizer, prompt_text):
    prompt_text = (args.padding_text if args.padding_text else PADDING_TEXT) + prompt_text
    return prompt_text


PREPROCESSING_FUNCTIONS = {
    "ctrl": prepare_ctrl_input,
    "xlm": prepare_xlm_input,
    "xlnet": prepare_xlnet_input,
    "transfo-xl": prepare_transfoxl_input,
}


def adjust_length_to_model(length, max_sequence_length):
    if length < 0 and max_sequence_length > 0:
        length = max_sequence_length
    elif 0 < max_sequence_length < length:
        length = max_sequence_length  # No generation bigger than model size
    elif length < 0:
        length = MAX_LENGTH  # avoid infinite loop
    return length

def collate_batch(zhutis,examples,device):

    attention_mask, position_id, input_ids = get_mask_and_position(zhutis)
    attention_mask=attention_mask.to(device)
    position_id = position_id.to(device)
    input_ids = input_ids.to(device)
    f_attention_mask=torch.ones(zhutis.shape[-1],input_ids.shape[-1])
    f_attention_mask = f_attention_mask.to(device)
    return {"input_ids": examples,  "zhutis": input_ids, "e_attention_mask": attention_mask,
            "e_position_ids": position_id, "f_attention_mask": f_attention_mask}

def get_mask_and_position(ids):
    # 输入topic信息而不是单独的tokenembedding,顿号544,句号545
    length=ids.shape[-1]
    num_sep=0
    for i in range(length):
        if ids[0][i].item() == 50259:
            num_sep+=1
    length_withoutsep=length-num_sep
    #id groups
    id_groups=[]
    #token groups
    token_groups=[]
    # print('local_rank')
    # print(local_rank)
    mask=torch.zeros([length_withoutsep,length_withoutsep])
    position=torch.arange(0, length_withoutsep, dtype=torch.long)

    #to get relation_position encoding
    relative_position=0
    # to determine whether come to the last token
    end_flag=True

    seps_now=0
    #[cls]
    id_groups.append(torch.tensor([0]))
    token_groups.append(ids.index_select(-1, torch.tensor([0])))
    mask[0, :] = 1.0
    mask[:, 0] = 1.0

    for i in range(length):
        if ids[0][i].item() == 50259:
            seps_now+=1
            ids_before=[]
            ids_convert=[]
            head=i
            for tail in range(i+1,length):
                if ids[0][tail].item() == 50259:
                    end_flag=False
                    break
                ids_convert.append(tail - seps_now)
                ids_before.append(tail)
            id_groups.append(torch.tensor(ids_convert))
            token_groups.append(ids.index_select(-1, torch.tensor(ids_before)))
            if end_flag==False:
                mask[head+1-seps_now:tail-seps_now,head+1-seps_now:tail-seps_now]=1.0
            else:
                mask[head+1-seps_now:tail+1-seps_now, head+1-seps_now:tail+1-seps_now] = 1.0

            # relative position record starts again
            relative_position = 1
            end_flag=True
        else:
            position[i-seps_now]=relative_position
            relative_position+=1
    new_input_ids=torch.cat((token_groups),-1)
    mask=mask.unsqueeze(0)
    position = position.unsqueeze(0)
    return mask,position,new_input_ids

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_type",
        default=None,
        type=str,
        required=True,
        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
    )
    parser.add_argument(
        "--model_name_or_path",
        default=None,
        type=str,
        required=True,
        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=20)
    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=1)
    parser.add_argument("--p", type=float, default=0.9)

    parser.add_argument("--padding_text", type=str, default="", help="Padding text for Transfo-XL and XLNet.")
    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("--no_cuda", action="store_true", help="Avoid using CUDA when available")
    parser.add_argument("--num_return_sequences", type=int, default=1, help="The number of samples to generate.")
    parser.add_argument("--input_file", type=str, default='inference_test.txt', help="input_file.")
    parser.add_argument("--concept_file", type=str, default='test_concept.txt', help="input_file.")
    parser.add_argument("--output_file", type=str, default='decode_result.txt', help="output_file.")



    args = parser.parse_args()

    args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
    args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()

    set_seed(args)

    # 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)
    model = model_class.from_pretrained(args.model_name_or_path)
    total_params = sum(p.numel() for p in model.parameters())
    print(f'{total_params:,} total parameters.')
    model.to(args.device)


    #args.length = adjust_length_to_model(args.length, max_sequence_length=model.config.max_position_embeddings)
    logger.info(args)

    results=[]
    with open(args.input_file,'r',encoding='utf-8') as f:
        prompts=f.readlines()
    with open(args.concept_file,'r',encoding='utf-8') as f:
        concepts=f.readlines()
    time=0
    old_text=''
    starter,ender=torch.cuda.Event(enable_timing=True),torch.cuda.Event(enable_timing=True)
    timings=np.zeros((len(prompts),1))

    # to test bart model, open these codes and open teg env and close 283-287
    #model = BartForConditionalGeneration.from_pretrained("facebook/bart-large")
    #tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
    #model.to(args.device)
    # model = T5ForConditionalGeneration.from_pretrained("t5-base")
    # tokenizer = T5Tokenizer.from_pretrained("t5-base")
    # model.to(args.device)
    for prompt_text,concept in zip(prompts,concepts):

        if prompt_text != old_text:
            old_text=prompt_text
            prompt_text = prompt_text.strip()
            # 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)
                encoded_prompt = tokenizer.encode(
                    preprocessed_prompt_text, add_special_tokens=False, return_tensors="pt",
                    add_space_before_punct_symbol=True
                )
            else:

                encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
            concept_prompt=tokenizer.encode(concept, add_special_tokens=False, return_tensors="pt")
            encoded_prompt = encoded_prompt.to(args.device)
            dic = collate_batch(concept_prompt, encoded_prompt, args.device)
            # flops, params = profile(model, inputs=(dic["input_ids"],dic["zhutis"]))
            # print(flops / 1e9, params / 1e6)  # flops单位G，para单位M
            #concept_prompt = concept_prompt.to(args.device)

            # for key in dic:
            #     print(key + ':' )
            #     print(dic[key])
            starter.record()
            print(dic["input_ids"])
            output_sequences = model.generate(
                input_ids=dic["input_ids"],
                # eos_token_id=58,
                num_return_sequences=1,
                max_length=20 + len(encoded_prompt[0]),
                temperature=1.0,
                top_k=1,
                top_p=0.9,
                repetition_penalty=1.0,
                length_penalty=1,
                do_sample=False,
                num_beams=5,
                #no_repeat_ngram_size=3,
                use_cache=None,
                zhutis=dic["zhutis"],
                e_attention_mask=dic["e_attention_mask"],
                e_position_ids=dic["e_position_ids"],
                f_attention_mask=dic["f_attention_mask"],
            )
            ender.record()
            torch.cuda.synchronize()
            curr_time=starter.elapsed_time(ender)
            timings[time]=curr_time
            print(curr_time)
            # Remove the batch dimension when returning multiple sequences
            if len(output_sequences.shape) > 2:
                output_sequences.squeeze_()

            for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
                print("=== GENERATED SEQUENCE {} ===".format(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]
                #print(text)
                # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
                total_sequence = (
                        text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True))+1:] + '\n'
                )

                results.append(total_sequence)
                print(prompt_text)
                print(tokenizer.encode(prompt_text))
                print(total_sequence)
            time = time + 1

        else:
            results.append(total_sequence)
            # print(prompt_text)
            # print(total_sequence)
    mean_syn=np.sum(timings)/time
    print(mean_syn)
    with open(args.output_file, 'w', encoding='utf-8') as f:
        f.writelines(results)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# 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 logging

import numpy as np
import torch

from transformers import (
    CTRLLMHeadModel,
    CTRLTokenizer,
    GPT2LMHeadModel,
    GPT2Tokenizer,
    OpenAIGPTLMHeadModel,
    OpenAIGPTTokenizer,
    TransfoXLLMHeadModel,
    TransfoXLTokenizer,
    XLMTokenizer,
    XLMWithLMHeadModel,
    XLNetLMHeadModel,
    XLNetTokenizer,
)
from thop import profile
from torchstat import stat
from transformers import BartForConditionalGeneration, BartTokenizer
from transformers import T5ForConditionalGeneration, T5Tokenizer



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 = {
    "gpt2": (GPT2LMHeadModel, GPT2Tokenizer),
    "ctrl": (CTRLLMHeadModel, CTRLTokenizer),
    "openai-gpt": (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
    "xlnet": (XLNetLMHeadModel, XLNetTokenizer),
    "transfo-xl": (TransfoXLLMHeadModel, TransfoXLTokenizer),
    "xlm": (XLMWithLMHeadModel, XLMTokenizer),
}

# Padding text to help Transformer-XL and XLNet with short prompts as proposed by <NAME>
# in https://github.com/rusiaaman/XLNet-gen#methodology
# and https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e
PADDING_TEXT = """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 <NAME>, narrates the
remainder of the story. 1883 Western Siberia,
a young <NAME> 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>"""


def set_seed(args):
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


#
# 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).")

    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")
    return prompt_text


def prepare_xlm_input(args, model, tokenizer, prompt_text):
    # kwargs = {"language": None, "mask_token_id": None}

    # Set the language
    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
        else:
            language = None
            while language not in available_languages:
                language = input("Using XLM. Select language in " + str(list(available_languages)) + " >>> ")

        model.config.lang_id = model.config.lang2id[language]
        # kwargs["language"] = tokenizer.lang2id[language]

    # TODO fix mask_token_id setup when configurations will be synchronized between models and tokenizers
    # XLM masked-language modeling (MLM) models need masked token
    # is_xlm_mlm = "mlm" in args.model_name_or_path
    # if is_xlm_mlm:
    #     kwargs["mask_token_id"] = tokenizer.mask_token_id

    return prompt_text


def prepare_xlnet_input(args, _, tokenizer, prompt_text):
    prompt_text = (args.padding_text if args.padding_text else PADDING_TEXT) + prompt_text
    return prompt_text


def prepare_transfoxl_input(args, _, tokenizer, prompt_text):
    prompt_text = (args.padding_text if args.padding_text else PADDING_TEXT) + prompt_text
    return prompt_text


PREPROCESSING_FUNCTIONS = {
    "ctrl": prepare_ctrl_input,
    "xlm": prepare_xlm_input,
    "xlnet": prepare_xlnet_input,
    "transfo-xl": prepare_transfoxl_input,
}


def adjust_length_to_model(length, max_sequence_length):
    if length < 0 and max_sequence_length > 0:
        length = max_sequence_length
    elif 0 < max_sequence_length < length:
        length = max_sequence_length  # No generation bigger than model size
    elif length < 0:
        length = MAX_LENGTH  # avoid infinite loop
    return length

def collate_batch(zhutis,examples,device):

    attention_mask, position_id, input_ids = get_mask_and_position(zhutis)
    attention_mask=attention_mask.to(device)
    position_id = position_id.to(device)
    input_ids = input_ids.to(device)
    f_attention_mask=torch.ones(zhutis.shape[-1],input_ids.shape[-1])
    f_attention_mask = f_attention_mask.to(device)
    return {"input_ids": examples,  "zhutis": input_ids, "e_attention_mask": attention_mask,
            "e_position_ids": position_id, "f_attention_mask": f_attention_mask}

def get_mask_and_position(ids):
    # 输入topic信息而不是单独的tokenembedding,顿号544,句号545
    length=ids.shape[-1]
    num_sep=0
    for i in range(length):
        if ids[0][i].item() == 50259:
            num_sep+=1
    length_withoutsep=length-num_sep
    #id groups
    id_groups=[]
    #token groups
    token_groups=[]
    # print('local_rank')
    # print(local_rank)
    mask=torch.zeros([length_withoutsep,length_withoutsep])
    position=torch.arange(0, length_withoutsep, dtype=torch.long)

    #to get relation_position encoding
    relative_position=0
    # to determine whether come to the last token
    end_flag=True

    seps_now=0
    #[cls]
    id_groups.append(torch.tensor([0]))
    token_groups.append(ids.index_select(-1, torch.tensor([0])))
    mask[0, :] = 1.0
    mask[:, 0] = 1.0

    for i in range(length):
        if ids[0][i].item() == 50259:
            seps_now+=1
            ids_before=[]
            ids_convert=[]
            head=i
            for tail in range(i+1,length):
                if ids[0][tail].item() == 50259:
                    end_flag=False
                    break
                ids_convert.append(tail - seps_now)
                ids_before.append(tail)
            id_groups.append(torch.tensor(ids_convert))
            token_groups.append(ids.index_select(-1, torch.tensor(ids_before)))
            if end_flag==False:
                mask[head+1-seps_now:tail-seps_now,head+1-seps_now:tail-seps_now]=1.0
            else:
                mask[head+1-seps_now:tail+1-seps_now, head+1-seps_now:tail+1-seps_now] = 1.0

            # relative position record starts again
            relative_position = 1
            end_flag=True
        else:
            position[i-seps_now]=relative_position
            relative_position+=1
    new_input_ids=torch.cat((token_groups),-1)
    mask=mask.unsqueeze(0)
    position = position.unsqueeze(0)
    return mask,position,new_input_ids

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_type",
        default=None,
        type=str,
        required=True,
        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
    )
    parser.add_argument(
        "--model_name_or_path",
        default=None,
        type=str,
        required=True,
        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=20)
    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=1)
    parser.add_argument("--p", type=float, default=0.9)

    parser.add_argument("--padding_text", type=str, default="", help="Padding text for Transfo-XL and XLNet.")
    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("--no_cuda", action="store_true", help="Avoid using CUDA when available")
    parser.add_argument("--num_return_sequences", type=int, default=1, help="The number of samples to generate.")
    parser.add_argument("--input_file", type=str, default='inference_test.txt', help="input_file.")
    parser.add_argument("--concept_file", type=str, default='test_concept.txt', help="input_file.")
    parser.add_argument("--output_file", type=str, default='decode_result.txt', help="output_file.")



    args = parser.parse_args()

    args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
    args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()

    set_seed(args)

    # 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)
    model = model_class.from_pretrained(args.model_name_or_path)
    total_params = sum(p.numel() for p in model.parameters())
    print(f'{total_params:,} total parameters.')
    model.to(args.device)


    #args.length = adjust_length_to_model(args.length, max_sequence_length=model.config.max_position_embeddings)
    logger.info(args)

    results=[]
    with open(args.input_file,'r',encoding='utf-8') as f:
        prompts=f.readlines()
    with open(args.concept_file,'r',encoding='utf-8') as f:
        concepts=f.readlines()
    time=0
    old_text=''
    starter,ender=torch.cuda.Event(enable_timing=True),torch.cuda.Event(enable_timing=True)
    timings=np.zeros((len(prompts),1))

    # to test bart model, open these codes and open teg env and close 283-287
    #model = BartForConditionalGeneration.from_pretrained("facebook/bart-large")
    #tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
    #model.to(args.device)
    # model = T5ForConditionalGeneration.from_pretrained("t5-base")
    # tokenizer = T5Tokenizer.from_pretrained("t5-base")
    # model.to(args.device)
    for prompt_text,concept in zip(prompts,concepts):

        if prompt_text != old_text:
            old_text=prompt_text
            prompt_text = prompt_text.strip()
            # 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)
                encoded_prompt = tokenizer.encode(
                    preprocessed_prompt_text, add_special_tokens=False, return_tensors="pt",
                    add_space_before_punct_symbol=True
                )
            else:

                encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
            concept_prompt=tokenizer.encode(concept, add_special_tokens=False, return_tensors="pt")
            encoded_prompt = encoded_prompt.to(args.device)
            dic = collate_batch(concept_prompt, encoded_prompt, args.device)
            # flops, params = profile(model, inputs=(dic["input_ids"],dic["zhutis"]))
            # print(flops / 1e9, params / 1e6)  # flops单位G，para单位M
            #concept_prompt = concept_prompt.to(args.device)

            # for key in dic:
            #     print(key + ':' )
            #     print(dic[key])
            starter.record()
            print(dic["input_ids"])
            output_sequences = model.generate(
                input_ids=dic["input_ids"],
                # eos_token_id=58,
                num_return_sequences=1,
                max_length=20 + len(encoded_prompt[0]),
                temperature=1.0,
                top_k=1,
                top_p=0.9,
                repetition_penalty=1.0,
                length_penalty=1,
                do_sample=False,
                num_beams=5,
                #no_repeat_ngram_size=3,
                use_cache=None,
                zhutis=dic["zhutis"],
                e_attention_mask=dic["e_attention_mask"],
                e_position_ids=dic["e_position_ids"],
                f_attention_mask=dic["f_attention_mask"],
            )
            ender.record()
            torch.cuda.synchronize()
            curr_time=starter.elapsed_time(ender)
            timings[time]=curr_time
            print(curr_time)
            # Remove the batch dimension when returning multiple sequences
            if len(output_sequences.shape) > 2:
                output_sequences.squeeze_()

            for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
                print("=== GENERATED SEQUENCE {} ===".format(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]
                #print(text)
                # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
                total_sequence = (
                        text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True))+1:] + '\n'
                )

                results.append(total_sequence)
                print(prompt_text)
                print(tokenizer.encode(prompt_text))
                print(total_sequence)
            time = time + 1

        else:
            results.append(total_sequence)
            # print(prompt_text)
            # print(total_sequence)
    mean_syn=np.sum(timings)/time
    print(mean_syn)
    with open(args.output_file, 'w', encoding='utf-8') as f:
        f.writelines(results)


if __name__ == "__main__":
    main()
#!/usr/bin/env python3 # 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) # Hardcoded max length to avoid infinite loop # Padding text to help Transformer-XL and XLNet with short prompts as proposed by <NAME> # in https://github.com/rusiaaman/XLNet-gen#methodology # and https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e 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 <NAME>, narrates the remainder of the story. 1883 Western Siberia, a young <NAME> 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> # # Functions to prepare models' input # # kwargs = {"language": None, "mask_token_id": None} # Set the language # kwargs["language"] = tokenizer.lang2id[language] # TODO fix mask_token_id setup when configurations will be synchronized between models and tokenizers # XLM masked-language modeling (MLM) models need masked token # is_xlm_mlm = "mlm" in args.model_name_or_path # if is_xlm_mlm: #     kwargs["mask_token_id"] = tokenizer.mask_token_id # No generation bigger than model size # avoid infinite loop # 输入topic信息而不是单独的tokenembedding,顿号544,句号545 #id groups #token groups # print('local_rank') # print(local_rank) #to get relation_position encoding # to determine whether come to the last token #[cls] # relative position record starts again # Initialize the model and tokenizer #args.length = adjust_length_to_model(args.length, max_sequence_length=model.config.max_position_embeddings) # to test bart model, open these codes and open teg env and close 283-287 #model = BartForConditionalGeneration.from_pretrained("facebook/bart-large") #tokenizer = BartTokenizer.from_pretrained("facebook/bart-large") #model.to(args.device) # model = T5ForConditionalGeneration.from_pretrained("t5-base") # tokenizer = T5Tokenizer.from_pretrained("t5-base") # model.to(args.device) # Different models need different input formatting and/or extra arguments # flops, params = profile(model, inputs=(dic["input_ids"],dic["zhutis"])) # print(flops / 1e9, params / 1e6)  # flops单位G，para单位M #concept_prompt = concept_prompt.to(args.device) # for key in dic: #     print(key + ':' ) #     print(dic[key]) # eos_token_id=58, #no_repeat_ngram_size=3, # Remove the batch dimension when returning multiple sequences # Decode text # Remove all text after the stop token #print(text) # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing # print(prompt_text) # print(total_sequence)