app.py

import gradio as gr
from cijiang import CiJiangRhymer
from cijiang.utils import print_results

rhymer = CiJiangRhymer(strict=False, tone=True)

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_results,
    inputs=[
        gr.Textbox(
            label="Input Text",
            placeholder="Enter your text with [M] markers for rhyme positions",
            lines=3,
            value="[M][M][M][M]"
        ),
        gr.Textbox(
            label="Target Rhyme",
            placeholder="Enter target rhyme pattern (Chinese characters or pinyin)",
            value="乱比高低"
        ),
        gr.Slider(
            minimum=1,
            maximum=50,
            value=20,
            step=1,
            label="Beam Width"
        ),
        gr.Slider(
            minimum=100,
            maximum=10000,
            value=4000,
            step=100,
            label="Number of Candidates"
        ),
        gr.Slider(
            minimum=1,
            maximum=20,
            value=10,
            step=1,
            label="Top Results to Show"
        )
    ],
    outputs=gr.Textbox(
        label="Generated Rhymes",
        lines=15,
        max_lines=20
    ),
    title="CiJiang Chinese Poetry Rhymer",
    description="Generate Chinese poetry rhymes using the CiJiang library. Use [M] markers to indicate positions where rhymes should be generated. Supports both Chinese characters and pinyin as target rhymes.",
    examples=[
        ["[M][M][M][M]", "乱比高低", 20, 4000, 10],
        ["[M][M][M][M]", "luan4 bi3 gao1 di1", 20, 4000, 10],
        ["[M][M][M]", "收视率", 20, 4000, 10],
        ["我总是天亮了[M][M]", "排队", 20, 4000, 10],
        ["我们两个人的verse价值一辆[M][M][M][M]", "韩国低级", 20, 4000, 10],
        ["报名嘻[M][M]", "八个", 20, 4000, 10],
        ["我还是觉得音乐的[M][M][M][M][M]", "背靠背钻研", 20, 4000, 10],
        ["我还是觉得音乐的[M][M][M][M][M]", "背靠背钻研", 300, 4000, 10]
    ]
)
demo.launch()