Update app.py

app.py

@@ -9,8 +9,8 @@ from transformers import pipeline
 @st.cache(allow_output_mutation=True)
 def load_model():
   model_ckpt = "bankholdup/rugpt3_song_writer"
-  tokenizer = GPT2Tokenizer.from_pretrained(model_ckpt,from_flax=True)
-  model = GPT2LMHeadModel.from_pretrained(model_ckpt,from_flax=True)
+  tokenizer = GPT2Tokenizer.from_pretrained(model_ckpt)
+  model = GPT2LMHeadModel.from_pretrained(model_ckpt)
   return tokenizer, model
 
 def set_seed(args):
@@ -23,24 +23,43 @@ def set_seed(args):
 
 title = st.title("Loading model")
 tokenizer, model = load_model()
-text_generation = pipeline("text-generation", model=model, tokenizer=tokenizer)
 title.title("ruGPT3 Song Writer")
 context = st.text_input("Введите начало песни", "Как дела? Как дела? Это новый кадиллак")
 
 if st.button("Поехали", help="Может занять какое-то время"):
     st.title(f"{context}")
     prefix_text = f"{context}"
-    generated_song = text_generation(prefix_text, max_length=200,         temperature=0.9, k=50, do_sample=True)[0]
-    for count, line in enumerate(generated_song['generated_text'].split("\n")):
-      if"<EOS>" in line:
-        break
-      if count == 0:
-        st.markdown(f"**{line[line.find('['):]}**")
-        continue
-      if "<BOS>" in line:
-        st.write(line[5:])
-        continue
-      if line.startswith("["):
-        st.markdown(f"**{line}**")
-        continue
-      st.write(line)
+    encoded_prompt = tokenizer.encode(prefix_text, add_special_tokens=False, return_tensors="pt")
+    output_sequences = model.generate(
+            input_ids=encoded_prompt,
+            max_length=200 + len(encoded_prompt[0]),
+            temperature=0.95,
+            top_k=50,
+            top_p=0.95,
+            repetition_penalty=1.0,
+            do_sample=True,
+            num_return_sequences=1,
+        )
+
+    # Remove the batch dimension when returning multiple sequences
+    if len(output_sequences.shape) > 2:
+        output_sequences.squeeze_()
+    
+    for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
+        print("ruGPT:".format(generated_sequence_idx + 1))
+        generated_sequence = generated_sequence.tolist()
+    
+        # Decode text
+        text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
+    
+        # Remove all text after the stop token
+        text = text[: text.find(args.stop_token) if args.stop_token else None]
+    
+        # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
+        total_sequence = (
+            prompt_text + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
+        )
+    
+        generated_sequences.append(total_sequence)
+        # os.system('clear')
+        st.write(total_sequence)