Update app.py

app.py

@@ -1,7 +1,45 @@
 import gradio as gr
+from cijiang import CiJiangRhymer
+from cijiang.utils import print_results
 
-def greet(name):
-    return "Hello " + name + "!!"
+rhymer = CiJiangRhymer(strict=False, tone=True)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+def get_results(input_text, target_rhyme, beam_width=20, num_candidates=4000):
+    out = rhymer.get_rhymes(text=input_text, target_rhyme=target_rhyme, beam_width=beam_width, num_candidates=num_candidates)
+    mask_count = text.count("[M]")
+    context = text.split('[M]')[0]
+
+    all_outputs = f"======= 韵脚: |{target_rhyme}|\n"
+    for i, (seq, log_prob) in enumerate(out[:top_results]): 
+        rhymes = seq[-mask_count:].split()
+        colored_rhymes = [Fore.RED + part + Style.RESET_ALL if idx < mask_count else part for idx, part in enumerate(rhymes)]
+        colored_rhymes = ''.join(colored_rhymes)  # Join the parts back together
+
+        all_outputs += f"{i+1}. {context}{colored_rhymes} (score: {log_prob:.3f})\n"
+    all_outputs += "=" * 40
+    return all_out
+
+# rhymer = CiJiangRhymer(strict=False, tone=True)
+
+# # 使用韵脚的汉字押韵
+# print_results(rhymer, text="[M][M][M][M]", target_rhyme="乱比高低")
+
+# # 使用韵脚的拼音押韵
+# print_results(rhymer, text="[M][M][M][M]", target_rhyme="luan4 bi3 gao1 di1")
+
+# # 允许声调不同
+# rhymer.tone = False
+# print_results(rhymer, text="[M][M][M]", target_rhyme="收视率")
+# rhymer.tone = True
+
+# # 上下文押韵
+# print_results(rhymer, text="我总是天亮了[M][M]", target_rhyme="排队")
+# print_results(rhymer, text="我们两个人的verse价值一辆[M][M][M][M]", target_rhyme="韩国低级")
+
+# # # 上下文非词汇押韵
+# print_results(rhymer, text="报名嘻[M][M]", target_rhyme="八个")
+# # 长韵脚需要更广的搜索空间
+# print_results(rhymer, text="我还是觉得音乐的[M][M][M][M][M]", target_rhyme="背靠背钻研")
+# print_results(rhymer, text="我还是觉得音乐的[M][M][M][M][M]", target_rhyme="背靠背钻研", beam_width=300)
+demo = gr.Interface(fn=get, inputs="text", outputs="text")
 demo.launch()