Update app.py

app.py

@@ -7,8 +7,8 @@ from transformers import GPT2Tokenizer, GPT2LMHeadModel
 
 @st.cache(allow_output_mutation=True)
 def load_model():
-  model_ckpt = "bankholdup/rugpt3_song_writer"
-#  model_ckpt = "bankholdup/mgpt_song_write"
+#  model_ckpt = "bankholdup/rugpt3_song_writer"
+  model_ckpt = "bankholdup/mgpt_song_write"
   tokenizer = GPT2Tokenizer.from_pretrained(model_ckpt)
   model = GPT2LMHeadModel.from_pretrained(model_ckpt)
   return tokenizer, model
@@ -32,7 +32,8 @@ if st.button("Поехали", help="Может занять какое-то в
         #st.write("temperature = {}".format(temperature))
         #st.write("_____________")
         prompt_text = f"{context}"
-        encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
+        #encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
+        encoded_prompt = tokenizer.encode(text, return_tensors="pt").cuda(device)
         output_sequences = model.generate(
                 input_ids=encoded_prompt,
                 max_length=200 + len(encoded_prompt[0]),
@@ -42,7 +43,10 @@ if st.button("Поехали", help="Может занять какое-то в
                 repetition_penalty=1.0,
                 do_sample=True,
                 num_return_sequences=1,
+                eos_token_id=5, 
+                pad_token=1,
             )
+            
         if len(output_sequences.shape) > 2:
             output_sequences.squeeze_()