# from transformers import BartTokenizer, BartForConditionalGeneration

# tokenizer = BartTokenizer.from_pretrained("facebook/mbart-large-cc25")
# TXT = "My dog is so <mask>."
# model = BartForConditionalGeneration.from_pretrained("facebook/mbart-large-cc25")
# input_ids = tokenizer.encode_plus(TXT, return_tensors='pt')['input_ids']  # batch_encode_plus([TXT]
# logits = model(input_ids, output_hidden_states=True)[0]
# masked_index = (input_ids[0] == tokenizer.mask_token_id).nonzero().item()
# probs = logits[0, masked_index].softmax(dim=0)
# values, predictions = probs.topk(5)
# print(tokenizer.decode(predictions).split())

# bart code
from transformers import BartForConditionalGeneration, BartTokenizer
# article_en = "<s> Durant was selected as the <mask> pick in the 2007 NBA draft by the Seattle <mask> after Greg Oden was selected by the <mask> with the first pick of the draft."
# article_en = "A contrastive learning method generally <mask> two components: finding positive samples and <mask> samples for each anchor sample; and building up an <mask> to discriminate them"
# model = BartForConditionalGeneration.from_pretrained("facebook/bart-large").cuda()
# tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")

# tokenizer.src_lang = "en_XX"
# encoded_en = tokenizer(article_en, return_tensors="pt")
# for key, value in encoded_en.items():
#     encoded_en[key] = value.cuda()
# print("input_ids: ", encoded_en['input_ids'])
# generated_tokens = model.generate(
#     encoded_en['input_ids'], 
#     forced_bos_token_id=tokenizer.bos_token_id,
#     max_length=100
# )
# result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
# # result: A contrastive learning method generally consists of two components: finding positive samples and negative samples for each anchor sample; and building up an image of the anchor samples to discriminate them


# mbart code
lang_code="zh_CN"
article_zh="<s><mask>合同已<mask>不到一年<mask>适当增加组合的持续时间。<mask>5月的出口增长率达到13个。请在使用前进行检查,选择信用差价过高的<mask>债券和未来可能的信用利差。截至<mask>基于<mask>报告期末资产净值比例的前十大资产支持证券投资细分主要与信贷债券收益率有关<mask>的变化,5.01%;该期间的利润为该期间的已<mask>该期间公<mask>价值变动的收益。<mask>,曲线的陡<mask>期<mask>"
import torch
from transformers import MBartForConditionalGeneration, MBartTokenizer
model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-cc25").cuda()
tokenizer = MBartTokenizer.from_pretrained("facebook/mbart-large-cc25")

tokenizer.src_lang = lang_code
tokenizer.tgt_lang = lang_code
encoded_en = tokenizer(article_zh, return_tensors="pt")
input_ids = encoded_en['input_ids']
input_ids = input_ids.tolist()[0]
new_input_ids = []
for i in range(len(input_ids)):
    item = input_ids[i]
    if item != 6:
        new_input_ids.append(item)
new_input_ids = torch.LongTensor([new_input_ids])
new_input_ids = new_input_ids.to(encoded_en['input_ids']).cuda()
for key, value in encoded_en.items():
    encoded_en[key] = value.cuda()
print(encoded_en['input_ids'])
print("new_input_ids: ", new_input_ids)
generated_tokens = model.generate(
    new_input_ids, 
    decoder_start_token_id=tokenizer.lang_code_to_id[lang_code],
    forced_bos_token_id=tokenizer.bos_token_id,
    max_length=300,
    num_beams=5
) # decoder_start_token_id is crutial
result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
print(result)
# result: contrastive learning method generally consists of two components: finding positive samples and negative samples for each anchor sample; and building up an algorithm to discriminate them
