Update app.py

app.py

@@ -31,6 +31,7 @@ def get_model():
     #model = AutoModelForCausalLM.from_pretrained("BigSalmon/InformalToFormalLincoln63Paraphrase")
     tokenizer = AutoTokenizer.from_pretrained("BigSalmon/InformalToFormalLincoln71Paraphrase")
     model = AutoModelForCausalLM.from_pretrained("BigSalmon/InformalToFormalLincoln71Paraphrase")
+    tokenizer2 = AutoTokenizer.from_pretrained("sberbank-ai/mGPT")
     model2 = AutoModelForCausalLM.from_pretrained("sberbank-ai/mGPT")
     return model, model2, tokenizer
     
@@ -72,13 +73,13 @@ def run_generate(text, bad_words):
   
 def run_generate2(text, bad_words):
   yo = []
-  input_ids = tokenizer.encode(text, return_tensors='pt')
-  res = len(tokenizer.encode(text))
+  input_ids = tokenizer2.encode(text, return_tensors='pt')
+  res = len(tokenizer2.encode(text))
   bad_words = bad_words.split()
   bad_word_ids = []
   for bad_word in bad_words: 
     bad_word = " " + bad_word
-    ids = tokenizer(bad_word).input_ids
+    ids = tokenizer2(bad_word).input_ids
     bad_word_ids.append(ids)
   sample_outputs = model2.generate(
     input_ids,
@@ -91,7 +92,7 @@ def run_generate2(text, bad_words):
     bad_words_ids=bad_word_ids
   )
   for i in range(number_of_outputs):
-    e = tokenizer.decode(sample_outputs[i])
+    e = tokenizer2.decode(sample_outputs[i])
     e = e.replace(text, "")
     yo.append(e)
   return yo
@@ -126,12 +127,12 @@ with st.form(key='my_form'):
     if submit_button4:
         text2 = str(text)
         print(text2)
-        text3 = tokenizer.encode(text2)
+        text3 = tokenizer2.encode(text2)
         myinput, past_key_values = torch.tensor([text3]), None
         myinput = myinput
         logits, past_key_values = model2(myinput, past_key_values = past_key_values, return_dict=False)
         logits = logits[0,-1]
         probabilities = torch.nn.functional.softmax(logits)
         best_logits, best_indices = logits.topk(logs_outputs)
-        best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]      
+        best_words = [tokenizer2.decode([idx.item()]) for idx in best_indices]      
         st.write(best_words)