Update app.py

app.py

@@ -27,9 +27,10 @@ def load_model():
     encoder = model.get_encoder()
     decoder = model.model.decoder
     lm_head = model.lm_head
-    return encoder, decoder, lm_head, tokenizer
+    start_token = model.config.decoder_start_token_id
+    return encoder, decoder, lm_head, tokenizer, start_token
 
-encoder, decoder, lm_head, tokenizer = load_model()
+encoder, decoder, lm_head, tokenizer, start_token = load_model()
 
 # Title
 st.title("Czech punctuation and capitalization restoration (BART-Small)")
@@ -56,7 +57,7 @@ if submit_button:
         max_length = 50
 
         # Empty input for generated text
-        generated_ids = torch.tensor([[model.config.decoder_start_token_id]])
+        generated_ids = torch.tensor([[start_token]])
 
         output_placeholder = st.empty()