streaming

app.py

@@ -2,13 +2,16 @@ import gradio as gr
 from nltk.tokenize.treebank import TreebankWordDetokenizer
 from somajo import SoMaJo
 
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM, TextStreamer, TextIteratorStreamer
+from threading import Thread
 from datasets import Dataset
 from transformers.pipelines.pt_utils import KeyDataset
 from hybrid_textnorm.lexicon import Lexicon
 from hybrid_textnorm.normalization import predict_type_normalization, reranked_normalization, prior_normalization
 from hybrid_textnorm.preprocess import recombine_tokens, german_transliterate
 from tqdm import tqdm
+import re
+from collections import Counter
 
 text_tokenizer = SoMaJo("de_CMC", split_camel_case=True)
 lexicon_dataset_name = 'aehrm/dtaec-lexicon'
@@ -20,78 +23,72 @@ def predict(input_str, model_name, progress=gr.Progress()):
     tokenized_sentences = list(text_tokenizer.tokenize_text([input_str]))
 
     if model_name == 'type normalizer':
-        output_sentences = predict_only_type_transformer(tokenized_sentences, progress)
+        stream = predict_only_type_transformer(tokenized_sentences, progress)
     elif model_name == 'type normalizer + lm':
-        output_sentences = predict_type_transformer_with_lm(tokenized_sentences, progress)
+        stream = predict_type_transformer_with_lm(tokenized_sentences, progress)
     elif model_name == 'transnormer':
-        output_sentences = predict_transnormer(tokenized_sentences, progress)
+        stream = predict_transnormer(tokenized_sentences, progress)
 
-    if type(output_sentences[0]) == list:
-        return "\n".join([detok.detokenize(recombine_tokens(sent)) for sent in output_sentences])
-    else:
-        return "\n".join(output_sentences)
+    accumulated = ""
+    for out in stream:
+        accumulated += out
+        yield accumulated
 
 def predict_transnormer(tokenized_sentences, progress):
     model_name = 'ybracke/transnormer-19c-beta-v02'
 
-    progress(0, desc='loading model')
-    pipe = pipeline(model=model_name)
+    #progress(0, desc='loading model')
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    streamer = TextIteratorStreamer(tokenizer)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
     raw_sentences = []
     for tokenized_sent in tokenized_sentences:
         sent = ''.join(tok.text + (' ' if tok.space_after else '') for tok in tokenized_sent)
-        raw_sentences.append(sent)
 
+        inputs = tokenizer([sent], return_tensors='pt')
 
-    progress(0, desc='running normalization')
-    ds = KeyDataset(Dataset.from_dict(dict(types=list(raw_sentences))), "types")
-
-    output_sentences = []
-    for out_sentence in progress.tqdm(pipe(ds, num_beams=1, max_length=1024)):
-        output_sentences.append(out_sentence[0]['generated_text'])
-
-    return output_sentences
+        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1000, num_beams=1)
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+        for new_text in streamer:
+            yield re.sub(r'(<pad>|</s>)', '', new_text)
+        yield '\n'
 
 
 
 def predict_only_type_transformer(tokenized_sentences, progress):
     type_model_name = 'aehrm/dtaec-type-normalizer'
-    progress(0, desc='loading model')
-
-    pipe = pipeline('text2text-generation', type_model_name)
+    #progress(0, desc='loading model')
+    type_model_tokenizer = AutoTokenizer.from_pretrained(type_model_name)
+    type_model = AutoModelForSeq2SeqLM.from_pretrained(type_model_name)
 
     transliterated_sentences = []
     for sentence in tokenized_sentences:
-        transliterated_sentences.append([german_transliterate(tok.text) for tok in sentence])
+        transliterated = [german_transliterate(tok.text) for tok in sentence]
 
-    oov_types = set(tok for sent in transliterated_sentences for tok in sent) - train_lexicon.keys()
-    oov_normalizations = {}
+        oov_replacement_probabilities = {}
+        oov_types = set(transliterated) - train_lexicon.keys() - oov_replacement_probabilities.keys()
+        #print('oov:', oov_types)
+        for input_type, probas in predict_type_normalization(oov_types, type_model_tokenizer, type_model, batch_size=8):
+            oov_replacement_probabilities[input_type] = probas
 
-    progress(0, desc='running normalization')
-    ds = KeyDataset(Dataset.from_dict(dict(types=list(oov_types))), "types")
-    for in_type, out in zip(ds, progress.tqdm(pipe(ds))):
-        oov_normalizations[in_type] = out[0]['generated_text']
-    
-    output_sentences = []
-    for sent in transliterated_sentences:
         output_sent = []
-        for t in sent:
+        for t in transliterated:
             if t in train_lexicon.keys():
                 output_sent.append(train_lexicon[t].most_common(1)[0][0])
-            elif t in oov_normalizations.keys():
-                output_sent.append(oov_normalizations[t])
+            elif t in oov_replacement_probabilities.keys():
+                output_sent.append(Counter(dict(oov_replacement_probabilities[t])).most_common(1)[0][0])
             else:
                 raise ValueError()
 
-        output_sentences.append(output_sent)
-
-    return output_sentences
+        yield detok.detokenize(recombine_tokens(output_sent)) + '\n'
 
 def predict_type_transformer_with_lm(tokenized_sentences, progress):
     type_model_name = 'aehrm/dtaec-type-normalizer'
     language_model_name = 'dbmdz/german-gpt2'
 
-    progress(0, desc='loading model')
+    #progress(0, desc='loading model')
     type_model_tokenizer = AutoTokenizer.from_pretrained(type_model_name)
     type_model = AutoModelForSeq2SeqLM.from_pretrained(type_model_name)
     language_model_tokenizer = AutoTokenizer.from_pretrained(language_model_name)
@@ -99,25 +96,18 @@ def predict_type_transformer_with_lm(tokenized_sentences, progress):
     if 'pad_token' not in language_model_tokenizer.special_tokens_map:
         language_model_tokenizer.add_special_tokens({'pad_token': '<pad>'})
 
-    transliterated_sentences = []
-    for sentence in tokenized_sentences:
-        transliterated_sentences.append([german_transliterate(tok.text) for tok in sentence])
-
-    oov_types = set(tok for sent in transliterated_sentences for tok in sent) - train_lexicon.keys()
     oov_replacement_probabilities = {}
+    for sentence in tokenized_sentences:
+        transliterated = [german_transliterate(tok.text) for tok in sentence]
+        oov_types = set(transliterated) - train_lexicon.keys() - oov_replacement_probabilities.keys()
 
-    progress(0, desc='running normalization')
-    for input_type, probas in progress.tqdm(predict_type_normalization(oov_types, type_model_tokenizer, type_model, batch_size=8), total=len(oov_types)):
-        oov_replacement_probabilities[input_type] = probas
+        #print('oov:', oov_types)
+        for input_type, probas in predict_type_normalization(oov_types, type_model_tokenizer, type_model, batch_size=8):
+            oov_replacement_probabilities[input_type] = probas
 
-    output_sentences = []
-    progress(0, desc='running LM re-ranking')
-    for hist_sent in progress.tqdm(transliterated_sentences):
-        predictions = reranked_normalization(hist_sent, train_lexicon, oov_replacement_probabilities, language_model_tokenizer, language_model, batch_size=1)
+        predictions = reranked_normalization(transliterated, train_lexicon, oov_replacement_probabilities, language_model_tokenizer, language_model, batch_size=1)
         best_pred, _, _, _ = predictions[0]
-        output_sentences.append(best_pred)
-
-    return output_sentences
+        yield detok.detokenize(recombine_tokens(best_pred)) + '\n'
     
 
 gradio_app = gr.Interface(