import gradio as gr from time import time from bertviz import model_view, head_view from bertviz_gradio import head_view_mod import faiss import torch import os # import nltk import argparse import random import numpy as np import pandas as pd from argparse import Namespace from tqdm.notebook import tqdm from torch.utils.data import DataLoader from functools import partial from transformers import AutoTokenizer, MarianTokenizer, AutoModel, AutoModelForSeq2SeqLM, MarianMTModel model_es = "Helsinki-NLP/opus-mt-en-es" model_fr = "Helsinki-NLP/opus-mt-en-fr" model_zh = "Helsinki-NLP/opus-mt-en-zh" model_sw = "Helsinki-NLP/opus-mt-en-sw" tokenizer_es = AutoTokenizer.from_pretrained(model_es) tokenizer_fr = AutoTokenizer.from_pretrained(model_fr) tokenizer_zh = AutoTokenizer.from_pretrained(model_zh) tokenizer_sw = AutoTokenizer.from_pretrained(model_sw) model_tr_es = MarianMTModel.from_pretrained(model_es) model_tr_fr = MarianMTModel.from_pretrained(model_fr) model_tr_zh = MarianMTModel.from_pretrained(model_zh) model_tr_sw = MarianMTModel.from_pretrained(model_sw) from faiss import write_index, read_index import pickle def load_index(model): with open('index/'+ model + '_metadata_ref.pkl', 'rb') as f: loaded_dict = pickle.load(f) for type in ['tokens','words']: for kind in ['input', 'output']: ## save index file name = 'index/'+ model + "_" + kind + "_"+ type + ".index" loaded_dict[kind][type][1] = read_index(name) # write_index(metadata_all[kind][type][1], name) return loaded_dict dict_models = { 'en-es': model_es, 'en-fr': model_fr, 'en-zh': model_zh, 'en-sw': model_sw, } dict_models_tr = { 'en-es': model_tr_es, 'en-fr': model_tr_fr, 'en-zh': model_tr_zh, 'en-sw': model_tr_sw, } dict_tokenizer_tr = { 'en-es': tokenizer_es, 'en-fr': tokenizer_fr, 'en-zh': tokenizer_zh, 'en-sw': tokenizer_sw, } dict_reference_faiss = { 'en-es': load_index('en-es'), } saliency_examples = [ "Peace of Mind: Protection for consumers.", "The sustainable development goals report: towards a rescue plan for people and planet", "We will leave no stone unturned to hold those responsible to account.", "The clock is now ticking on our work to finalise the remaining key legislative proposals presented by this Commission to ensure that citizens and businesses can reap the benefits of our policy actions.", "Pumpkins, squash and gourds, fresh or chilled, excluding courgettes", "The labour market participation of mothers with infants has even deteriorated over the past two decades, often impacting their career and incomes for years.", ] contrastive_examples = [ ["Peace of Mind: Protection for consumers.", "Paz mental: protección de los consumidores", "Paz de la mente: protección de los consumidores"], ["the slaughterer has finished his work.", "l'abatteur a terminé son travail.", "l'abatteuse a terminé son travail."], ['A fundamental shift is needed - in commitment, solidarity, financing and action - to put the world on a better path.', '需要在承诺、团结、筹资和行动方面进行根本转变,使世界走上更美好的道路。', '我们需要从根本上转变承诺、团结、资助和行动，使世界走上更美好的道路。',] ] #Load challenge set examples df_challenge_set = pd.read_csv("challenge_sets.csv") arr_challenge_set = df_challenge_set.values arr_challenge_set = [[x[2], x[3], x[4], x[5]] for x in arr_challenge_set] def get_k_prob_tokens(transition_scores, result, model, k_values=5): tokenizer_tr = dict_tokenizer_tr[model] gen_sequences = result.sequences[:, 1:] result_output = [] # First beam only... bs = 0 text = ' ' for tok, score, i_step in zip(gen_sequences[bs], transition_scores[bs],range(len(gen_sequences[bs]))): 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 def split_token_from_sequences(sequences, model) -> dict : n_sentences = len(sequences) 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. 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 }] id_dict_pos = {} step_i = 0 cont = True words_by_step = [] #[['bos' for i in range(n_sentences)]] while cont: # append to dict_parent for all beams of step_i cont = False step_words = [] for beam in range(n_sentences): app_text = '' if step_i < len(gen_sequences_texts[beam]): app_text = gen_sequences_texts[beam][step_i] cont = True step_words.append(app_text) words_by_step.append(step_words) print(words_by_step) for i_bs, step_w in enumerate(step_words): if not step_w in ['', '']: #new id if the same word is not in another beam (?) [beam[i] was a token id] #parent id = previous word and previous step. # new_parent_id = "-".join([str(beam[i]) for i in range(step_i)]) new_id = "-".join([str(words_by_step[i][i_bs])+ '-' + str(i) for i in range(step_i+1)]) parent_id = "-".join([words_by_step[i][i_bs] + '-' + str(i) for i in range(step_i) ]) # new_id = step_w +'-' + str(step_i) # parent_id = words_by_step[step_i-1][i_bs] + '-' + str(step_i -1) next_word_flag = 1 if step_i == 0 : parent_id = 'bos--1' ## if the dict already exists remove it, if it is not a root... ## root?? then next is '' else: next_word_flag = len(gen_sequences_texts[i_bs][step_i]) > step_i ## Not in step_i = 0; if next_word_flag: 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 else: 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 ## Tokenization def compute_tokenization(inputs, targets, w1, model): colors = ['tok-first-color', 'tok-second-color', 'tok-third-color', 'tok-fourth-color'] len_colors = len(colors); inputs = inputs.input_ids html_tokens = "" i = 0 for sentence in inputs: html_tokens += "

" # print("TOKENS", inputs, targets) # print("input", [dict_tokenizer_tr[model].decode(tok) for tok in sentence]) tokens = [dict_tokenizer_tr[model].decode(tok) for tok in sentence] for token in tokens: token = token.replace("<", "<'") # .substring(0, token.length - 2) html_tokens += "" + token + " " i +=1 html_tokens += "

" i = 0 # for tgt_sentence in targets : html_tokens_tgt = "" html_tokens_tgt += "

" # print("targets", [dict_tokenizer_tr[model].decode(tok) for tok in targets]) # print("targets", dict_tokenizer_tr[model].decode(targets)) tokens = [dict_tokenizer_tr[model].decode(tok) for tok in targets] for token in tokens: token = token.replace("<", "<'") # .substring(0, token.length - 2) html_tokens_tgt += "" + token + " " i +=1 html_tokens_tgt += "

" # print("HTML", html_tokens, html_tokens_tgt) return html_tokens, html_tokens_tgt def create_vocab_multiple(embeddings_list, model): """_summary_ Args: embeddings_list (list): embedding array Returns: Dict: vocabulary of tokens' embeddings """ print("START VOCAB CREATION MULTIPLE \n \n ") vocab = {} ## add embedds. sentence_tokens_text_list = [] for embeddings in embeddings_list: tokens_id = embeddings['tokens'] # [[tokens_id]x n_sentences ] for sent_i, sentence in enumerate(tokens_id): sentence_tokens = [] for tok_i, token in enumerate(sentence): sentence_tokens.append(token) if not (token in vocab): vocab[token] = { 'token' : token, 'count': 1, # 'text': embeddings['texts'][sent_i][tok_i], 'text': dict_tokenizer_tr[model].decode([token]), # 'text': src_token_lists[sent_i][tok_i], 'embed': embeddings['embeddings'][sent_i][tok_i]} else: vocab[token]['count'] = vocab[token]['count'] + 1 # print(vocab) sentence_tokens_text_list.append(sentence_tokens) print("END VOCAB CREATION MULTIPLE \n \n ") return vocab, sentence_tokens_text_list def vocab_words_all_prefix(token_embeddings, model, sufix="@@",prefix = '▁' ): vocab = {} # inf_model = dict_models_tr[model] sentence_words_text_list = [] if prefix : n_prefix = len(prefix) for input_sentences in token_embeddings: # n_tokens_in_word for sent_i, sentence in enumerate(input_sentences['tokens']): words_text_list = [] # embedding = input_sentences['embed'][sent_i] word = '' tokens_ids = [] embeddings = [] ids_to_tokens = dict_tokenizer_tr[model].convert_ids_to_tokens(sentence) # print("validate same len", len(sentence) == len(ids_to_tokens), len(sentence), len(ids_to_tokens), ids_to_tokens) to_save= False for tok_i, token_text in enumerate(ids_to_tokens): token_id = sentence[tok_i] if token_text[:n_prefix] == prefix : #first we save the previous word if to_save: vocab[word] = { 'word' : word, 'text': word, 'count': 1, 'tokens_ids' : tokens_ids, 'embed': np.mean(np.array(embeddings), 0).tolist() } words_text_list.append(word) #word is starting if prefix tokens_ids = [token_id] embeddings = [input_sentences['embeddings'][sent_i][tok_i]] word = token_text[n_prefix:] ## if word to_save = True else : if (token_text in dict_tokenizer_tr[model].special_tokens_map.values()): # print('final or save', token_text, token_id, to_save, word) if to_save: # vocab[word] = ids vocab[word] = { 'word' : word, 'text': word, 'count': 1, 'tokens_ids' : tokens_ids, 'embed': np.mean(np.array(embeddings), 0).tolist() } words_text_list.append(word) #special token is one token element, no continuation # vocab[token_text] = [token_id] tokens_ids = [token_id] embeddings = [input_sentences['embeddings'][sent_i][tok_i]] vocab[token_text] = { 'word' : token_text, 'count': 1, 'text': word, 'tokens_ids' : tokens_ids, 'embed': np.mean(np.array(embeddings), 0).tolist() } words_text_list.append(token_text) to_save = False else: # is a continuation; we do not know if it is final; we don't save here. to_save = True word += token_text tokens_ids.append(token_id) embeddings.append(input_sentences['embeddings'][sent_i][tok_i]) if to_save: # print('final save', token_text, token_id, to_save, word) vocab[word] = tokens_ids if not (word in vocab): vocab[word] = { 'word' : word, 'count': 1, 'text': word, 'tokens_ids' : tokens_ids, 'embed': np.mean(np.array(embeddings), 0).tolist() } words_text_list.append(word) else: vocab[word]['count'] = vocab[word]['count'] + 1 sentence_words_text_list.append(words_text_list) return vocab, sentence_words_text_list def search_query_vocab(index, vocab_queries, topk = 10, limited_search = []): """ the embed queries are a vocabulary of words : embds_input_voc Args: index (_type_): faiss index embed_queries (_type_): vocab format. { 'token' : token, 'count': 1, 'text': src_token_lists[sent_i][tok_i], 'embed': embeddings[0]['embeddings'][sent_i][tok_i] } nb_ids (_type_): hash to find the token_id w.r.t the faiss index id. topk (int, optional): nb of similar tokens. Defaults to 10. Returns: _type_: Distance matrix D, indices matrix I and tokens ids (using nb_ids) """ # nb_qi_ids = [] ##ordered ids list nb_q_embds = [] ##ordered embeddings list metadata = {} qi_pos = 0 for key , token_values in vocab_queries.items(): # nb_qi_ids.append(token_values['token']) # for x in vocab_tokens] metadata[qi_pos] = {'word': token_values['word'], 'tokens': token_values['tokens_ids'], 'text': token_values['text']} qi_pos += 1 nb_q_embds.append(token_values['embed']) # for x in vocab_tokens] xq = np.array(nb_q_embds).astype('float32') #elements to query D,I = index.search(xq, topk) return D,I, metadata def search_query_vocab_token(index, vocab_queries, topk = 10, limited_search = []): """ the embed queries are a vocabulary of words : embds_input_vov Returns: _type_: Distance matrix D, indices matrix I and tokens ids (using nb_ids) """ # nb_qi_ids = [] ##ordered ids list nb_q_embds = [] ##ordered embeddings list metadata = {} qi_pos = 0 for key , token_values in vocab_queries.items(): # nb_qi_ids.append(token_values['token']) # for x in vocab_tokens] metadata[qi_pos] = {'token': token_values['token'], 'text': token_values['text']} qi_pos += 1 nb_q_embds.append(token_values['embed']) # for x in vocab_tokens] xq = np.array(nb_q_embds).astype('float32') #elements to query D,I = index.search(xq, topk) return D,I, metadata def build_search(query_embeddings, model,type="input"): metadata_all = dict_reference_faiss[model] # ## biuld vocab for index vocab_queries, sentence_tokens_list = create_vocab_multiple(query_embeddings, model) words_vocab_queries, sentence_words_list = vocab_words_all_prefix(query_embeddings, model, sufix="@@",prefix="▁") index_vor_tokens = metadata_all[type]['tokens'][1] md_tokens = metadata_all[type]['tokens'][2] D, I, meta = search_query_vocab_token(index_vor_tokens, vocab_queries) qi_pos = 0 similar_tokens = {} # similar_tokens = [] for dist, ind in zip(D,I): try: # similar_tokens.append({ similar_tokens[str(meta[qi_pos]['token'])] = { 'token': meta[qi_pos]['token'], 'text': meta[qi_pos]['text'], # 'text': dict_tokenizer_tr[model].decode(meta[qi_pos]['token']) # 'text': meta[qi_pos]['text'], "similar_topk": [md_tokens[i_index]['token'] for i_index in ind if (i_index != -1) ], "distance": [dist[i] for (i, i_index) in enumerate(ind) if (i_index != -1)], } # ) except: print("\n ERROR ", qi_pos, dist, ind) qi_pos += 1 index_vor_words = metadata_all[type]['words'][1] md_words = metadata_all[type]['words'][2] Dw, Iw, metaw = search_query_vocab(index_vor_words, words_vocab_queries) # D, I, meta, vocab_words, sentence_words_list = result_input['words']# [2] # D ; I ; meta qi_pos = 0 # similar_words = [] similar_words = {} for dist, ind in zip(Dw,Iw): try: # similar_words.append({ similar_words[str(metaw[qi_pos]['word']) ] = { 'word': metaw[qi_pos]['word'], 'text': metaw[qi_pos]['word'], "similar_topk": [md_words[i_index]['word'] for i_index in ind if (i_index != -1) ], "distance": [dist[i] for (i, i_index) in enumerate(ind) if (i_index != -1)], } # ) except: print("\n ERROR ", qi_pos, dist, ind) qi_pos += 1 return {'tokens': {'D': D, 'I': I, 'meta': meta, 'vocab_queries': vocab_queries, 'similar':similar_tokens, 'sentence_key_list': sentence_tokens_list}, 'words': {'D':Dw,'I': Iw, 'meta': metaw, 'vocab_queries':words_vocab_queries, 'sentence_key_list': sentence_words_list, 'similar': similar_words} } from sklearn.manifold import TSNE def embds_input_projection_vocab(vocab, key="token"): t0 = time() nb_ids = [] ##ordered ids list nb_embds = [] ##ordered embeddings list nb_text = [] ##ordered embeddings list tnse_error = [] for _ , token_values in vocab.items(): tnse_error.append([0,0]) nb_ids.append(token_values[key]) # for x in vocab_tokens] nb_text.append(token_values['text']) # for x in vocab_tokens] nb_embds.append(token_values['embed']) # for x in vocab_tokens] X = np.array(nb_embds).astype('float32') #elements to project try: tsne = TSNE(random_state=0, n_iter=1000) tsne_results = tsne.fit_transform(X) tsne_results = np.c_[tsne_results, nb_ids, nb_text, range(len(nb_ids))] ## creates a zip array : [[TNSE[X,Y], tokenid, token_text], ...] except: tsne_results = np.c_[tnse_error, nb_ids, nb_text, range(len(nb_ids))] ## creates a zip array : [[TNSE[X,Y], tokenid, token_text], ...] t1 = time() print("t-SNE: %.2g sec" % (t1 - t0)) # print(tsne_results) return tsne_results.tolist() def filtered_projection(similar_key, vocab, model, type="input", key="word"): metadata_all = dict_reference_faiss[model] vocab_proj = vocab.copy() ## tnse projection Input words source_words_voc_similar = set() # for words_set in similar_key: for key_i in similar_key: words_set = similar_key[key_i] source_words_voc_similar.update(words_set['similar_topk']) # print(len(source_words_voc_similar)) # source_embeddings_filtered = {key: metadata_all['input']['words'][0][key] for key in source_words_voc_similar} source_embeddings_filtered = {key_value: metadata_all[type][key][0][key_value] for key_value in source_words_voc_similar} vocab_proj.update(source_embeddings_filtered) ## vocab_proj add try: result_TSNE = embds_input_projection_vocab(vocab_proj, key=key[:-1]) ## singular => without 's' dict_projected_embds_all = {str(embds[2]): [embds[0], embds[1], embds[2], embds[3], embds[4]] for embds in result_TSNE} except: print('TSNE error', type, key) dict_projected_embds_all = {} # print(result_TSNE) return dict_projected_embds_all def get_bertvis_data(input_text, lg_model): tokenizer_tr = dict_tokenizer_tr[lg_model] model_tr = dict_models_tr[lg_model] # input_ids = tokenizer_tr(input_text, return_tensors="pt", padding=True) input_ids = tokenizer_tr(input_text, return_tensors="pt", padding=False) result_att = model_tr.generate(**input_ids, num_beams=4, num_return_sequences=4, return_dict_in_generate=True, output_attentions =True, output_scores=True, ) # tokenizer_tr.convert_ids_to_tokens(result_att.sequences[0]) # tokenizer_tr.convert_ids_to_tokens(input_ids.input_ids[0]) tgt_text = tokenizer_tr.decode(result_att.sequences[0], skip_special_tokens=True) outputs = model_tr(input_ids=input_ids.input_ids, decoder_input_ids=result_att.sequences[:1], output_attentions =True, ) html_attentions = head_view_mod( encoder_attention = outputs.encoder_attentions, cross_attention = outputs.cross_attentions, decoder_attention = outputs.decoder_attentions, encoder_tokens = tokenizer_tr.convert_ids_to_tokens(input_ids.input_ids[0]), decoder_tokens = tokenizer_tr.convert_ids_to_tokens(result_att.sequences[0]), html_action='gradio' ) return html_attentions, tgt_text, result_att, outputs def translation_model(w1, model): #translate and get internal values and visualizations; # src_text = saliency_examples[0] inputs = dict_tokenizer_tr[model](w1, return_tensors="pt", padding=True) num_ret_seq = 4 translated = dict_models_tr[model].generate(**inputs, num_beams=4, num_return_sequences=num_ret_seq, return_dict_in_generate=True, output_attentions =True, output_hidden_states = True, output_scores=True,) beam_dict = split_token_from_sequences(translated.sequences,model ) tgt_text = dict_tokenizer_tr[model].decode(translated.sequences[0], skip_special_tokens=True) ## Attentions outputs = dict_models_tr[model](input_ids=inputs.input_ids, decoder_input_ids=translated.sequences[:1], output_attentions =True, ) encoder_tokens = dict_tokenizer_tr[model].convert_ids_to_tokens(inputs.input_ids[0]) decoder_tokens = dict_tokenizer_tr[model].convert_ids_to_tokens(translated.sequences[0]) # decoder_tokens = [tok for tok in decoder_tokens if tok != ''] # decoder_tokens = [tok for tok in decoder_tokens if tok != ''] # html_attentions = head_view_mod( # encoder_attention = outputs.encoder_attentions, # cross_attention = outputs.cross_attentions, # decoder_attention = outputs.decoder_attentions, # encoder_tokens = encoder_tokens, # decoder_tokens = decoder_tokens, # html_action='gradio' # ) html_attentions_enc = head_view_mod( encoder_attention = outputs.encoder_attentions, encoder_tokens = encoder_tokens, decoder_tokens = decoder_tokens, html_action='gradio' ) html_attentions_dec = head_view_mod( # encoder_attention = outputs.encoder_attentions, decoder_attention = outputs.decoder_attentions, encoder_tokens = encoder_tokens, decoder_tokens = decoder_tokens, html_action='gradio' ) html_attentions_cross = head_view_mod( cross_attention = outputs.cross_attentions, encoder_tokens = encoder_tokens, decoder_tokens = decoder_tokens, html_action='gradio' ) # tokenization html_in, html_out = compute_tokenization(inputs, translated.sequences[0],w1, model) transition_scores = dict_models_tr[model].compute_transition_scores( translated.sequences, translated.scores, translated.beam_indices , normalize_logits=True ) prob_tokens = get_k_prob_tokens(transition_scores, translated, model, k_values=10) input_embeddings = dict_models_tr[model].get_encoder().embed_tokens(inputs.input_ids) target_embeddings = dict_models_tr[model].get_decoder().embed_tokens(translated.sequences) return [tgt_text, [beam_dict,prob_tokens, html_in, html_out, translated, inputs.input_ids,input_embeddings,target_embeddings], [html_attentions_enc['params'], html_attentions_enc['html2'].data], [html_attentions_dec['params'], html_attentions_dec['html2'].data], [html_attentions_cross['params'], html_attentions_cross['html2'].data] ] html = """

Exploring top-k probable tokens

... top 10 tokens generated at each step ...

Exploring the Beam Search sequence generation

""" html_tok = """
... tokenization visualization ...
""" html_embd = """
... token embeddings visualization ...
""" html_tok_target ="""
... tokenization visualization ...
""" html_embd_target= """
... token embeddings visualization ...
""" html_att_enc = """
... Encoder self attention only -- last layer and mean across heads ... Always read from left to right
""" html_att_cross = """
... Encoder-decoder cross attention only -- last layer and mean across heads ...
""" html_att_dec = """
... decoder self attention only -- last layer and mean across heads ...
""" def sentence_maker2(w1,j2): print(w1,j2) return "in sentence22..." def first_function(w1, model): global metadata_all #translate and get internal values sentences = w1.split("\n") all_sentences = [] translated_text = '' input_embeddings = [] output_embeddings = [] for sentence in sentences : # print(sentence, end=";") params = translation_model(sentence, model) all_sentences.append(params) # print(len(params)) translated_text += params[0] + ' \n' input_embeddings.append({ 'embeddings': params[1][6].detach(), ## create a vocabulary with the set of embeddings 'tokens': params[1][3+2].tolist(), # one translation = one sentence # 'texts' : dict_tokenizer_tr[model].decode(params[2].tolist()) }) output_embeddings.append({ 'embeddings' : params[1][7].detach(), 'tokens': params[1][3+1].sequences.tolist(), # 'texts' : dict_tokenizer_tr[model].decode(params[1].sequences.tolist()) }) ## load_reference; ## Build FAISS index # ---> preload faiss using the respective model with a initial dataset. ## dict_reference_faiss[model] = metadata_all [per language] # result_input = build_reference(input_embeddings,model) # result_output = build_reference(output_embeddings,model) # metadata_all = {'input': result_input, 'output': result_output} ## Build FAISS index # ---> preload faiss using the respective model with a initial dataset. result_search = {} result_search['input'] = build_search(input_embeddings, model, type='input') result_search['output'] = build_search(output_embeddings, model, type='output') json_out = {'input': {'tokens': {}, 'words': {}}, 'output': {'tokens': {}, 'words': {}}} dict_projected = {} for type in ['input', 'output']: dict_projected[type] = {} for key in ['tokens', 'words']: similar_key = result_search[type][key]['similar'] vocab = result_search[type][key]['vocab_queries'] dict_projected[type][key] = filtered_projection(similar_key, vocab, model, type=type, key=key) json_out[type][key]['similar_queries'] = similar_key json_out[type][key]['tnse'] = dict_projected[type][key] json_out[type][key]['key_text_list'] = result_search[type][key]['sentence_key_list'] ## bertviz # paramsbv, tgtbv = get_bertvis_data(w1, model) # params.append(json_out) html_att_enc = params[2][1]#.root_div_id = "bertviz_enc" html_att_dec = params[3][1] html_att_cross = params[4][1] params = [params[0], params[1], json_out, params[2][0], params[3][0], params[4][0]] # params.append([tgt, params['params'], params['html2'].data] return [translated_text, params, html_att_enc, html_att_dec, html_att_cross] def second_function(w1,j2): # json_value = {'one':1}# return f"{w1['two']} in sentence22..." # to transfer the data to json. print("second_function -- after the js", w1,j2) return "transition to second js function finished." with gr.Blocks(js="plotsjs.js") as demo: gr.Markdown( """ # MAKE NMT Workshop \t `Literacy task` """) gr.Markdown( """ ### Translation """) gr.Markdown( """ 1. Select the language pair for the translation """) radio_c = gr.Radio(choices=['en-zh', 'en-es', 'en-fr', 'en-sw'], value="en-es", label= '', container=False) gr.Markdown( """ 2. Source text to translate """) in_text = gr.Textbox(label="source text") with gr.Accordion("Optional: Challenge selection:", open=False): gr.Markdown( """ ### select an example from the challenge set listed bellow """) challenge_ex = gr.Textbox(label="Challenge", interactive=False) category_minor = gr.Textbox(label="category_minor", interactive=False) category_major = gr.Textbox(label="category_major", interactive=False) with gr.Accordion("Examples:"): gr.Examples(arr_challenge_set,[in_text, challenge_ex,category_minor,category_major], label="") btn = gr.Button("Translate") with gr.Accordion("3. Review the source tokenization:", open=False): input_tokenisation = gr.HTML(html_tok) with gr.Accordion("4. Review similar source tokens in the embedding space:", open=False): input_embd= gr.HTML(html_embd) with gr.Accordion("5. Review the attention between the source tokens:", open=False): gr.Markdown( """ `Bertviz ` """) input_embd= gr.HTML(html_att_enc) enc_html = gr.HTML() gr.Markdown( """ ### Text is translated into Target Language """) out_text = gr.Textbox(label="target text") with gr.Accordion("1. Review the target tokenization:", open=False): target_tokenisation = gr.HTML(html_tok_target) with gr.Accordion("2. Review similar target tokens in the embedding space:", open=False): target_embd= gr.HTML(html_embd_target) with gr.Accordion("3. Review the attention between the target and source tokens:", open=False): gr.Markdown( """ `Bertviz -cross attention` """) input_embd= gr.HTML(html_att_cross) cross_html = gr.HTML() with gr.Accordion("4. Review the attention between the target tokens:", open=False): gr.Markdown( """ `Bertviz -dec attention` """) input_embd= gr.HTML(html_att_dec) dec_html = gr.HTML() with gr.Accordion("6. Review the alternative translations tokens:", open=False): gr.Markdown( """ Generation process : `topk - beam search ` """) input_mic = gr.HTML(html) out_text2 = gr.Textbox(visible=False) var2 = gr.JSON(visible=False) btn.click(first_function, [in_text, radio_c], [out_text,var2,enc_html, dec_html, cross_html], js="(in_text,radio_c) => testFn_out(in_text,radio_c)") #should return an output comp. out_text.change(second_function, [out_text, var2], out_text2, js="(out_text,var2) => testFn_out_json(var2)") # # run script function on load, # demo.load(None,None,None,js="plotsjs.js") if __name__ == "__main__": demo.launch()