Spaces:

supermy
/

couplet

Runtime error

App Files Files Community

couplet / app.py

supermy

Update app.py

af532f0 over 1 year ago

raw history blame contribute delete

No virus

3.66 kB

	import torch

	import gradio as gr
	import torch.nn.functional as F

	from transformers import BertTokenizer, GPT2LMHeadModel

	tokenizer = BertTokenizer.from_pretrained("supermy/couplet-gpt2")
	model = GPT2LMHeadModel.from_pretrained("supermy/couplet-gpt2")
	model.eval()

	# top_k或top_p解码策略，仅保留top_k个或累积概率到达top_p的标记，其他标记设为filter_value，后续在选取标记的过程中会取不到值设为无穷小。
	# 从模型输出的logit里面，划分出概率最高的几个
	def top_k_top_p_filtering( logits, top_k=0, top_p=0.0, filter_value=-float('Inf') ):
	#确保输出logit维度为1行若干列的矩阵，便于处理
	assert logits.dim() == 1
	#将top_k的值初始化为logit元素个数和设定的top_k之间的最小值
	top_k = min( top_k, logits.size(-1) )
	if top_k > 0:
	# 去除掉概率值小于top_k里最后一个token概率的后续token
	# torch.topk()返回最后一维最大的top_k个元素，返回值为二维(values,indices)
	# ...表示其他维度由计算机自行推断
	indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
	logits[indices_to_remove] = filter_value # 对于topk之外的其他元素的logits值设为负无穷
	if top_p > 0.0:
	sorted_logits, sorted_indices = torch.sort(logits, descending=True) # 对logits进行递减排序
	cumulative_probs = torch.cumsum( F.softmax(sorted_logits, dim=-1), dim=-1 )
	sorted_indices_to_remove = cumulative_probs > top_p
	sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
	sorted_indices_to_remove[..., 0] = 0
	indices_to_remove = sorted_indices[sorted_indices_to_remove]
	logits[indices_to_remove] = filter_value
	return logits

	def generate(input_text):
	input_ids = [tokenizer.cls_token_id]
	input_ids.extend( tokenizer.encode(input_text + "-", add_special_tokens=False) )
	input_ids = torch.tensor( [input_ids] )

	# 最多生成max_len个token
	generated = []
	for _ in range(100):
	output = model(input_ids)
	# 对于已生成的结果generated中的每个token添加一个重复惩罚项，降低其生成概率
	next_token_logits = output.logits[0, -1, :]
	# 对于[UNK]的概率设为无穷小，也就是说模型的预测结果不可能是[UNK]这个token
	next_token_logits[ tokenizer.convert_tokens_to_ids('[UNK]') ] = -float('Inf')
	# 使用top_k_top_p_filtering函数，按照top_k和top_p的值，对预测结果进行筛选
	filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=8, top_p=1)
	# torch.multinomial表示从候选集合中无放回地进行抽取num_samples个元素，权重越高，抽到的几率越高，返回元素的下标
	next_token = torch.multinomial( F.softmax(filtered_logits, dim=-1), num_samples=1 )
	if next_token == tokenizer.sep_token_id:
	break
	generated.append( next_token.item() )
	input_ids = torch.cat( (input_ids, next_token.unsqueeze(0)), dim=1 )

	return "".join( tokenizer.convert_ids_to_tokens(generated) )

	if __name__ == "__main__":

	gr.Markdown("![visitor badge](https://visitor-badge.glitch.me/badge?page_id=supermy/couplet/)"),
	gr.Interface(
	fn=generate,
	inputs=gr.Textbox(lines=1, placeholder="请在此输入上联【居末尾带句号。】：燕子归来，问昔日雕梁何处。",value="燕子归来，问昔日雕梁何处。",label="上联"),
	outputs=gr.Textbox(lines=1, placeholder="此处显示生成的下联",label="下联"),
	).launch()