zh-decode and best probabilities

app.py

@@ -78,19 +78,18 @@ def get_k_prob_tokens(transition_scores, result, model, k_values=5):
 	gen_sequences = result.sequences[:, 1:]
 
 	result_output = []
-	# bs_alt = []
-	# bs_alt_scores = []
 
 	# First beam only... 
 	bs = 0
 	text = ' '
 	for tok, score, i_step in zip(gen_sequences[bs], transition_scores[bs],range(len(gen_sequences[bs]))):
-		# bs_alt.append([tokenizer_tr.decode(tok) for tok in result.scores[i_step][bs].topk(k_values).indices ] )
-		# bs_alt_scores.append(np.exp(result.scores[i_step][bs].topk(k_values).values))
 
-		bs_alt = [tokenizer_tr.decode(tok) for tok in result.scores[i_step][bs].topk(k_values).indices ]
-		bs_alt_scores = np.exp(result.scores[i_step][bs].topk(k_values).values)
-		result_output.append([np.array(result.scores[i_step][bs].topk(k_values).indices), np.array(bs_alt_scores),bs_alt])
+		beam_i = result.beam_indices[0][i_step]
+		if beam_i < 0:
+			beam_i = bs
+		bs_alt = [tokenizer_tr.decode(tok) for tok in result.scores[i_step][beam_i].topk(k_values).indices ]
+		bs_alt_scores = np.exp(result.scores[i_step][beam_i].topk(k_values).values)
+		result_output.append([np.array(result.scores[i_step][beam_i].topk(k_values).indices), np.array(bs_alt_scores),bs_alt])
 
 	return result_output
 
@@ -100,15 +99,19 @@ def split_token_from_sequences(sequences, model) -> dict :
 
 	gen_sequences_texts = []
 	for bs in range(n_sentences): 
+		# gen_sequences_texts.append(dict_tokenizer_tr[model].decode(sequences[:, 1:][bs],  skip_special_tokens=True).split(' '))
 		#### decoder per token.
-		gen_sequences_texts.append(dict_tokenizer_tr[model].decode(sequences[:, 1:][bs],  skip_special_tokens=True).split(' '))
-	print(gen_sequences_texts)
-	score = 0
+		seq_bs = []
+		
+		for token in sequences[:, 1:][bs]:
+			seq_bs.append(dict_tokenizer_tr[model].decode(token,  skip_special_tokens=True))
+		gen_sequences_texts.append(seq_bs)
 
+	score = 0
 	#raw dict is bos
 	text = 'bos'
 	new_id = text +'--1'
-	dict_parent = [{'id': new_id, 'parentId': None , 'text': text, 'name': 'bos', 'prob':score }]
+	dict_parent = [{'id': new_id, 'parentId': None , 'text': text, 'name': 'bos', 'prob': score }]
 	id_dict_pos = {}
 	step_i = 0
 	cont = True
@@ -151,8 +154,10 @@ def split_token_from_sequences(sequences, model) -> dict :
 						dict_parent.append({'id': new_id, 'parentId': parent_id , 'text': step_w, 'name': step_w, 'prob' : score })
 						id_dict_pos[new_id] = len(dict_parent) - 1
 				else: 
-					dict_parent.append({'id': new_id, 'parentId': parent_id , 'text': step_w, 'name': step_w, 'prob' : score })
-					id_dict_pos[new_id] = len(dict_parent) - 1
+					if not (new_id in id_dict_pos):
+						dict_parent.append({'id': new_id, 'parentId': parent_id , 'text': step_w, 'name': step_w, 'prob' : score  })
+						id_dict_pos[new_id] = len(dict_parent) - 1
+
 
 		step_i += 1
 	return dict_parent