Update app.py

app.py

@@ -135,8 +135,8 @@ def load_model(
         pretrained_model, cache_dir=cache_dir
     ).to(device)
     config = BertConfig.from_pretrained(bert, output_hidden_states=True)
-    bert_tokenizer: str = BertTokenizer.from_pretrained(bert, config=config)
-    bert_model: str = BertModel.from_pretrained(bert).to(device)
+    bert_tokenizer: str = BertTokenizer.from_pretrained(bert)
+    bert_model: str = BertModel.from_pretrained(bert,config=config).to(device)
     model.eval()
     bert_model.eval()
     return tokenizer, model, bert_tokenizer, bert_model
@@ -153,7 +153,6 @@ def find_algnments(
     bpe_target_map = []
     for i in translated_text.split():
         bpe_target_map += len(bert_tokenizer.tokenize(i)) * [i]
-    st.success(bert_model(**source_tokens))    
     source_embedding = bert_model(**source_tokens).hidden_states[8]
     target_embedding = bert_model(**target_tokens).hidden_states[8]
     target_embedding = target_embedding.transpose(-1, -2)