Update app.py

app.py

@@ -1,7 +1,17 @@
+import transformers
 import numpy as np
 import torch
-import gradio as gr
-from gradio import mix
+import streamlit as st
+
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+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)
+  return tokenizer, model
 
 def set_seed(args):
     rd = np.random.randint(100000)
@@ -11,23 +21,26 @@ def set_seed(args):
     if args.n_gpu > 0:
         torch.cuda.manual_seed_all(rd)
 
-title = "ruGPT3 Song Writer"
-description = "Generate russian songs via fine-tuned ruGPT3"
-
-io = gr.Interface.load("models/bankholdup/rugpt3_song_writer")
-
-examples = [
-    ['Как дела? Как дела? Это новый кадиллак']
-]
-
-def inference(text):
-    return io(text)
+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("Введите начало песни", "Как дела? Как дела? Это новый кадиллак")
 
-gr.Interface(
-    inference, 
-    [gr.inputs.Textbox(label="Input text")],
-    gr.outputs.Textbox(label="Output text"),
-    examples=examples,
-    title=title,
-    description=description,
-    ).launch(enable_queue=True,cache_examples=True)
+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)