Adapt to API updates

app.py

@@ -30,78 +30,117 @@ def process(query_type, corpus_desc, engine_desc, query, maxnum, request: gr.Req
         print(result)
     return result
 
+def format_tokenization_info(result):
+    if not ('token_ids' in result and 'tokens' in result):
+        return ''
+    token_ids = result['token_ids']
+    tokens = result['tokens']
+    t = '[' + " ".join(['"' + token.replace('Ġ', ' ') + '"' for token in tokens]) + '] ' + str(token_ids)
+    return t
+def format_tokenization_info_nested(result):
+    if not ('token_idsss' in result and 'tokensss' in result):
+        return ''
+    token_idsss = result['token_idsss']
+    tokensss = result['tokensss']
+    ttt = []
+    for token_idss, tokenss in zip(token_idsss, tokensss):
+        tt = []
+        for token_ids, tokens in zip(token_idss, tokenss):
+            t = '[' + " ".join(['"' + token.replace('Ġ', ' ') + '"' for token in tokens]) + '] ' + str(token_ids)
+            tt.append(t)
+        tt = '\n'.join(tt)
+        ttt.append(tt)
+    ttt = '\n\n'.join(ttt)
+    return ttt
+def format_doc(doc):
+    formatted = []
+    if doc['doc_len'] == doc['disp_len']:
+        header = f'[Document #{doc["doc_ix"]}, length = {doc["doc_len"]} tokens]\n\n'
+    else:
+        header = f'[Document #{doc["doc_ix"]}, length = {doc["doc_len"]} tokens ({doc["disp_len"]} tokens displayed)]\n\n'
+    formatted.append((header, None))
+    formatted += doc['spans']
+    return formatted
+
 def count(corpus_desc, engine_desc, query, request: gr.Request):
     result = process('count', corpus_desc, engine_desc, query, None, request)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
-    tokenized = '' if 'tokenized' not in result else result['tokenized']
+    tokenization_info = format_tokenization_info(result)
     if 'error' in result:
         count = result['error']
     else:
         count = f'{result["count"]:,}'
-    return latency, tokenized, count
+    return latency, tokenization_info, count
 
 def prob(corpus_desc, engine_desc, query, request: gr.Request):
     result = process('prob', corpus_desc, engine_desc, query, None, request)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
-    tokenized = '' if 'tokenized' not in result else result['tokenized']
+    tokenization_info = format_tokenization_info(result)
     if 'error' in result:
         prob = result['error']
     elif result['prompt_cnt'] == 0:
         prob = '(n-1)-gram is not found in the corpus'
     else:
         prob = f'{result["prob"]:.4f} ({result["cont_cnt"]:,} / {result["prompt_cnt"]:,})'
-    return latency, tokenized, prob
+    return latency, tokenization_info, prob
 
 def ntd(corpus_desc, engine_desc, query, request: gr.Request):
     result = process('ntd', corpus_desc, engine_desc, query, None, request)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
-    tokenized = '' if 'tokenized' not in result else result['tokenized']
+    tokenization_info = format_tokenization_info(result)
     if 'error' in result:
         ntd = result['error']
     else:
-        ntd = result['ntd']
+        result_by_token_id = result['result_by_token_id']
+        ntd = {}
+        for token_id, r in result_by_token_id.items():
+            ntd[f'{r["token"]} ({r["cont_cnt"]} / {result["prompt_cnt"]})'] = r['prob']
         if ntd == {}:
             ntd = '(n-1)-gram is not found in the corpus'
-    return latency, tokenized, ntd
+    return latency, tokenization_info, ntd
 
 def infgram_prob(corpus_desc, engine_desc, query, request: gr.Request):
     result = process('infgram_prob', corpus_desc, engine_desc, query, None, request)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
-    tokenized = '' if 'tokenized' not in result else result['tokenized']
+    tokenization_info = format_tokenization_info(result)
     if 'error' in result:
         longest_suffix = ''
         prob = result['error']
     else:
         longest_suffix = result['longest_suffix']
         prob = f'{result["prob"]:.4f} ({result["cont_cnt"]:,} / {result["prompt_cnt"]:,})'
-    return latency, tokenized, longest_suffix, prob
+    return latency, tokenization_info, longest_suffix, prob
 
 def infgram_ntd(corpus_desc, engine_desc, query, request: gr.Request):
     result = process('infgram_ntd', corpus_desc, engine_desc, query, None, request)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
-    tokenized = '' if 'tokenized' not in result else result['tokenized']
+    tokenization_info = format_tokenization_info(result)
     if 'error' in result:
         longest_suffix = ''
         ntd = result['error']
     else:
         longest_suffix = result['longest_suffix']
-        ntd = result['ntd']
-    return latency, tokenized, longest_suffix, ntd
+        result_by_token_id = result['result_by_token_id']
+        ntd = {}
+        for token_id, r in result_by_token_id.items():
+            ntd[f'{r["token"]} ({r["cont_cnt"]} / {result["prompt_cnt"]})'] = r['prob']
+    return latency, tokenization_info, longest_suffix, ntd
 
 def search_docs(corpus_desc, engine_desc, query, maxnum, request: gr.Request):
     result = process('search_docs', corpus_desc, engine_desc, query, maxnum, request)
     latency = '' if 'latency' not in result else f'{result["latency"]:.3f}'
-    tokenized = '' if 'tokenized' not in result else result['tokenized']
+    tokenization_info = format_tokenization_info_nested(result)
     if 'error' in result:
         message = result['error']
         docs = [[] for _ in range(10)]
     else:
         message = result['message']
-        docs = result['docs']
+        docs = result['documents']
+        docs = [format_doc(doc) for doc in docs]
     docs = docs[:maxnum]
     while len(docs) < 10:
         docs.append([])
-    return latency, tokenized, message, docs[0], docs[1], docs[2], docs[3], docs[4], docs[5], docs[6], docs[7], docs[8], docs[9]
+    return latency, tokenization_info, message, docs[0], docs[1], docs[2], docs[3], docs[4], docs[5], docs[6], docs[7], docs[8], docs[9]
 
 def analyze_document(corpus_desc, engine_desc, query, request: gr.Request):
     result = process('analyze_document', corpus_desc, engine_desc, query, None, request)