update

app.py

@@ -8,10 +8,13 @@ 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
 
 text_tokenizer = SoMaJo("de_CMC", split_camel_case=True)
 lexicon_dataset_name = 'aehrm/dtaec-lexicon'
 train_lexicon = Lexicon.from_dataset(lexicon_dataset_name, split='train')
+detok = TreebankWordDetokenizer()
+
 
 def predict(input_str, model_name, progress=gr.Progress()):
     tokenized_sentences = list(text_tokenizer.tokenize_text([input_str]))
@@ -24,7 +27,6 @@ def predict(input_str, model_name, progress=gr.Progress()):
         output_sentences = predict_transnormer(tokenized_sentences, progress)
 
     if type(output_sentences[0]) == list:
-        detok = TreebankWordDetokenizer()
         return "\n".join([detok.detokenize(recombine_tokens(sent)) for sent in output_sentences])
     else:
         return "\n".join(output_sentences)
@@ -32,18 +34,20 @@ def predict(input_str, model_name, progress=gr.Progress()):
 def predict_transnormer(tokenized_sentences, progress):
     model_name = 'ybracke/transnormer-19c-beta-v02'
 
-    progress(0, desc='running normalization')
-    pipe = pipeline(model='ybracke/transnormer-19c-beta-v02')
+    progress(0, desc='loading model')
+    pipe = pipeline(model=model_name)
 
     raw_sentences = []
     for tokenized_sent in tokenized_sentences:
-        raw_sentences.append(''.join(tok.text + (' ' if tok.space_after else '') for tok in tokenized_sent))
+        sent = ''.join(tok.text + (' ' if tok.space_after else '') for tok in tokenized_sent)
+        raw_sentences.append(sent)
+
 
     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=4, max_length=1000)):
+    for out_sentence in progress.tqdm(pipe(ds, num_beams=4, max_length=1024)):
         output_sentences.append(out_sentence[0]['generated_text'])
 
     return output_sentences
@@ -107,6 +111,7 @@ def predict_type_transformer_with_lm(tokenized_sentences, progress):
         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)
         best_pred, _, _, _ = predictions[0]