Update app.py

app.py

@@ -1,25 +1,61 @@
 import gradio as gr
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+import torch
+from random import sample
 
-def request_download(IMAGE_URL):
-    import requests
-    r = requests.get(IMAGE_URL)
-    with open('img.jpg', 'wb') as f:
-        f.write(r.content)   
-        
-
-demo = gr.Blocks()
-
-def change_img():
-    request_download('https://y.qq.com/music/photo_new/T002R300x300M0000007ujpB2dTrlW_2.jpg')
-    return gr.Image.update(value='img.jpg')
-
-with demo:
-    album_image = gr.Image(value='fws.jpg', visible=True)
-    fetch_songs = gr.Button(value="更新").style(full_width=True)
-    fetch_songs.click(
-        fn=change_img,
-        inputs=None,
-        outputs=album_image
-    )
-    
-demo.launch()
+model_path = 'souljoy/t5-chinese-lyric'
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+tokenizer = T5Tokenizer.from_pretrained(model_path)
+model = T5ForConditionalGeneration.from_pretrained(model_path).to(device)
+keywords_list = []
+with open('keywords', mode='r', encoding='utf-8') as f:
+    for line in f:
+        keywords_list.append(line.strip())
+
+
+def key_cng_change(v_list):
+    return '，'.join(v_list)
+
+
+def post_process(text):
+    return text.replace("\\n", "\n").replace("\\t", "\t")
+
+
+def send(my_singer_txt, my_song_txt, my_key_txt):
+    if len(my_singer_txt) == 0:
+        lyric = '请输入歌手名'
+    elif len(my_song_txt) == 0:
+        lyric = '请输入歌名'
+    elif len(my_key_txt) == 0:
+        lyric = '歌词主题词'
+    elif len(my_key_txt.split('，')) > 10:
+        lyric = '歌词主太多！最多10个'
+    else:
+        text = '用户：写一首歌，歌手“{}”，歌曲“{}”，以“{}”为主题。\\n小L：'.format(my_singer_txt, my_song_txt, my_key_txt)
+        encoding = tokenizer(text=[text], truncation=True, padding=True, max_length=64, return_tensors="pt").to(device)
+        out = model.generate(**encoding, return_dict_in_generate=True, output_scores=False, max_new_tokens=512,
+                             do_sample=True, top_p=1, temperature=0.7, no_repeat_ngram_size=3)
+        out_text = tokenizer.batch_decode(out["sequences"], skip_special_tokens=True)
+        lyric = post_process(out_text[0])
+
+    return lyric, gr.CheckboxGroup.update(choices=sample(keywords_list, 100), value=[])
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("""#### 歌词创作机器人 📻 🎵 """)
+    with gr.Row():
+        with gr.Column():
+            singer_txt = gr.Textbox(label='歌手名', placeholder='请输入你的名字，或者其他歌手名字')
+            song_txt = gr.Textbox(label='歌名', placeholder='请输入要创作的歌词名称')
+            key_txt = gr.Textbox(label='歌词主题词', placeholder='请输入主题词，以中文逗号（，）分割，或点击下方“主题词样例”。最多10个')
+            send_button = gr.Button(value="提交")
+            key_cng = gr.CheckboxGroup(label='主题词样例', choices=sample(keywords_list, 100))
+            
+
+        with gr.Column():
+            lyric_txt = gr.Textbox(label='生成歌词')
+    key_cng.change(key_cng_change, [key_cng], [key_txt])
+    send_button.click(send, [singer_txt, song_txt, key_txt], [lyric_txt,key_cng])
+
+if __name__ == "__main__":
+    demo.launch()